81 lines
3.6 KiB
Go
81 lines
3.6 KiB
Go
package controller
|
|
|
|
import (
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/QuantumNous/new-api/constant"
|
|
"github.com/QuantumNous/new-api/middleware"
|
|
"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/require"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"testing"
|
|
)
|
|
|
|
func TestListModelsAutomaticPoolAllFormats(t *testing.T) {
|
|
db := setupModelListControllerTestDB(t)
|
|
withSelfUseModeDisabled(t)
|
|
oldGroups, oldRatios, oldModelRatios := setting.UserUsableGroups2JSONString(), ratio_setting.GroupRatio2JSONString(), ratio_setting.ModelRatio2JSONString()
|
|
oldMemory := common.MemoryCacheEnabled
|
|
common.MemoryCacheEnabled = false
|
|
t.Cleanup(func() {
|
|
common.MemoryCacheEnabled = oldMemory
|
|
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(oldGroups))
|
|
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(oldRatios))
|
|
require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(oldModelRatios))
|
|
})
|
|
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"list-a":"","list-b":""}`))
|
|
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"list-a":1,"list-b":1}`))
|
|
require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{"zz-list-a":1,"zz-list-b":1,"zz-shared":1,"zz-disabled":1}`))
|
|
require.NoError(t, db.Create(&[]model.Channel{{Id: 5601, Type: constant.ChannelTypeOpenAI, Status: 1}, {Id: 5602, Type: constant.ChannelTypeOpenAI, Status: 1}, {Id: 5603, Status: 2}}).Error)
|
|
require.NoError(t, db.Create(&[]model.Ability{
|
|
{ChannelId: 5601, Group: "list-a", Model: "zz-list-a", Enabled: true},
|
|
{ChannelId: 5602, Group: "list-b", Model: "zz-list-b", Enabled: true},
|
|
{ChannelId: 5601, Group: "list-a", Model: "zz-shared", Enabled: true},
|
|
{ChannelId: 5602, Group: "list-b", Model: "zz-shared", Enabled: true},
|
|
{ChannelId: 5603, Group: "list-b", Model: "zz-disabled", Enabled: true},
|
|
{ChannelId: 5602, Group: "list-b", Model: "zz-unpriced", Enabled: true},
|
|
}).Error)
|
|
for _, preference := range []string{"cost", "speed", "stable"} {
|
|
for _, limited := range []bool{false, true} {
|
|
for _, format := range []int{constant.ChannelTypeOpenAI, constant.ChannelTypeAnthropic, constant.ChannelTypeGemini} {
|
|
policy, err := common.Marshal(model.TokenRoutingPolicy{Version: 1, Mode: "smart", AutoPreference: preference})
|
|
require.NoError(t, err)
|
|
token := &model.Token{Group: "auto", RoutingPolicy: string(policy), ModelLimitsEnabled: limited, ModelLimits: "zz-list-b,zz-disabled,zz-unpriced"}
|
|
engine := gin.New()
|
|
engine.GET("/v1/models", func(c *gin.Context) {
|
|
common.SetContextKey(c, constant.ContextKeyUserGroup, "list-a")
|
|
require.NoError(t, middleware.SetupContextForToken(c, token))
|
|
ListModels(c, format)
|
|
})
|
|
res := httptest.NewRecorder()
|
|
engine.ServeHTTP(res, httptest.NewRequest(http.MethodGet, "/v1/models", nil))
|
|
require.Equal(t, 200, res.Code)
|
|
var body struct {
|
|
Data []struct {
|
|
ID string `json:"id"`
|
|
} `json:"data"`
|
|
Models []struct {
|
|
Name string `json:"name"`
|
|
} `json:"models"`
|
|
}
|
|
require.NoError(t, common.Unmarshal(res.Body.Bytes(), &body))
|
|
names := []string{}
|
|
for _, item := range body.Data {
|
|
names = append(names, item.ID)
|
|
}
|
|
for _, item := range body.Models {
|
|
names = append(names, item.Name)
|
|
}
|
|
expected := []string{"zz-list-a", "zz-list-b", "zz-shared"}
|
|
if limited {
|
|
expected = []string{"zz-list-b"}
|
|
}
|
|
require.ElementsMatch(t, expected, names, "preference=%s limited=%v format=%d", preference, limited, format)
|
|
}
|
|
}
|
|
}
|
|
}
|