86 lines
3.5 KiB
Go
86 lines
3.5 KiB
Go
package controller
|
|
|
|
import (
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"testing"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/QuantumNous/new-api/constant"
|
|
"github.com/QuantumNous/new-api/model"
|
|
"github.com/QuantumNous/new-api/setting"
|
|
"github.com/QuantumNous/new-api/setting/ratio_setting"
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func TestSelectableGroupsRequireEnabledChannel(t *testing.T) {
|
|
originalUsableGroups := setting.UserUsableGroups2JSONString()
|
|
originalRatios := ratio_setting.GroupRatio2JSONString()
|
|
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"auto":"Auto","active":"Active","default":"Default","paused":"Paused","no-channel":"No channel"}`))
|
|
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"active":0.2,"default":1,"paused":0.3,"no-channel":0.4}`))
|
|
t.Cleanup(func() {
|
|
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(originalUsableGroups))
|
|
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(originalRatios))
|
|
model.InvalidatePricingCache()
|
|
})
|
|
|
|
db := setupModelListControllerTestDB(t)
|
|
require.NoError(t, db.Create(&model.User{
|
|
Id: 7001,
|
|
Username: "group-availability-user",
|
|
Password: "password",
|
|
Group: "default",
|
|
Status: common.UserStatusEnabled,
|
|
}).Error)
|
|
require.NoError(t, db.Create(&[]model.Channel{
|
|
{Id: 11, Type: constant.ChannelTypeOpenAI, Key: "active-key", Name: "active", Status: common.ChannelStatusEnabled},
|
|
{Id: 12, Type: constant.ChannelTypeOpenAI, Key: "default-key", Name: "default", Status: common.ChannelStatusManuallyDisabled},
|
|
{Id: 13, Type: constant.ChannelTypeOpenAI, Key: "paused-key", Name: "paused", Status: common.ChannelStatusAutoDisabled},
|
|
}).Error)
|
|
require.NoError(t, db.Create(&[]model.Ability{
|
|
{Group: "active", Model: "gpt-4o", ChannelId: 11, Enabled: true},
|
|
{Group: "default", Model: "gpt-4o", ChannelId: 12, Enabled: true},
|
|
{Group: "paused", Model: "gpt-4o", ChannelId: 13, Enabled: true},
|
|
}).Error)
|
|
model.InvalidatePricingCache()
|
|
|
|
userRecorder := httptest.NewRecorder()
|
|
userContext, _ := gin.CreateTestContext(userRecorder)
|
|
userContext.Request = httptest.NewRequest(http.MethodGet, "/api/user/groups", nil)
|
|
userContext.Set("id", 7001)
|
|
GetUserGroups(userContext)
|
|
|
|
var userResponse struct {
|
|
Success bool `json:"success"`
|
|
Data map[string]map[string]interface{} `json:"data"`
|
|
}
|
|
require.NoError(t, common.Unmarshal(userRecorder.Body.Bytes(), &userResponse))
|
|
require.True(t, userResponse.Success)
|
|
assert.ElementsMatch(t, []string{"active", "auto"}, mapKeys(userResponse.Data))
|
|
|
|
pricingRecorder := httptest.NewRecorder()
|
|
pricingContext, _ := gin.CreateTestContext(pricingRecorder)
|
|
pricingContext.Request = httptest.NewRequest(http.MethodGet, "/api/pricing", nil)
|
|
GetPricing(pricingContext)
|
|
|
|
var pricingResponse struct {
|
|
Success bool `json:"success"`
|
|
GroupRatio map[string]float64 `json:"group_ratio"`
|
|
UsableGroup map[string]string `json:"usable_group"`
|
|
}
|
|
require.NoError(t, common.Unmarshal(pricingRecorder.Body.Bytes(), &pricingResponse))
|
|
require.True(t, pricingResponse.Success)
|
|
assert.Equal(t, map[string]float64{"active": 0.2}, pricingResponse.GroupRatio)
|
|
assert.Equal(t, map[string]string{"active": "Active", "auto": "Auto"}, pricingResponse.UsableGroup)
|
|
}
|
|
|
|
func mapKeys[V any](values map[string]V) []string {
|
|
keys := make([]string, 0, len(values))
|
|
for key := range values {
|
|
keys = append(keys, key)
|
|
}
|
|
return keys
|
|
}
|