98 lines
4.4 KiB
Go
98 lines
4.4 KiB
Go
package service
|
|
|
|
import (
|
|
"fmt"
|
|
"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/glebarez/sqlite"
|
|
"github.com/stretchr/testify/require"
|
|
"gorm.io/driver/mysql"
|
|
"gorm.io/driver/postgres"
|
|
"gorm.io/gorm"
|
|
"gorm.io/gorm/schema"
|
|
"net/http/httptest"
|
|
"os"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestAutomaticPoolModelsDatabaseMatrix(t *testing.T) {
|
|
for _, dialect := range []string{"sqlite", "mysql", "postgres"} {
|
|
t.Run(dialect, func(t *testing.T) {
|
|
var driver gorm.Dialector
|
|
switch dialect {
|
|
case "sqlite":
|
|
driver = sqlite.Open(":memory:")
|
|
case "mysql":
|
|
if os.Getenv("TEST_MYSQL_DSN") == "" {
|
|
t.Skip("no isolated MySQL database")
|
|
}
|
|
driver = mysql.Open(os.Getenv("TEST_MYSQL_DSN"))
|
|
case "postgres":
|
|
if os.Getenv("TEST_POSTGRES_DSN") == "" {
|
|
t.Skip("no isolated PostgreSQL database")
|
|
}
|
|
driver = postgres.Open(os.Getenv("TEST_POSTGRES_DSN"))
|
|
}
|
|
db, err := gorm.Open(driver, &gorm.Config{NamingStrategy: schema.NamingStrategy{TablePrefix: fmt.Sprintf("automodel_%d_", time.Now().UnixNano())}})
|
|
require.NoError(t, err)
|
|
oldDB, oldRedis, oldMemory := model.DB, common.RedisEnabled, common.MemoryCacheEnabled
|
|
oldGroups, oldRatios := setting.UserUsableGroups2JSONString(), ratio_setting.GroupRatio2JSONString()
|
|
model.DB = db
|
|
common.RedisEnabled = false
|
|
common.MemoryCacheEnabled = false
|
|
require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.Ability{}))
|
|
t.Cleanup(func() {
|
|
require.NoError(t, db.Migrator().DropTable(&model.Ability{}, &model.Channel{}))
|
|
model.DB = oldDB
|
|
common.RedisEnabled = oldRedis
|
|
common.MemoryCacheEnabled = oldMemory
|
|
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(oldGroups))
|
|
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(oldRatios))
|
|
conn, e := db.DB()
|
|
require.NoError(t, e)
|
|
require.NoError(t, conn.Close())
|
|
})
|
|
var version string
|
|
query := "select version()"
|
|
if dialect == "sqlite" {
|
|
query = "select sqlite_version()"
|
|
}
|
|
require.NoError(t, db.Raw(query).Scan(&version).Error)
|
|
t.Log("database version:", version)
|
|
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"auto-a":"","auto-b":""}`))
|
|
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"auto-a":1,"auto-b":1,"private":1}`))
|
|
require.NoError(t, db.Create(&[]model.Channel{{Id: 561, Status: 1}, {Id: 562, Status: 1}, {Id: 563, Status: 2}, {Id: 564, Status: 1}}).Error)
|
|
require.NoError(t, db.Create(&[]model.Ability{
|
|
{ChannelId: 561, Group: "auto-a", Model: "model-a", Enabled: true},
|
|
{ChannelId: 561, Group: "auto-a", Model: "shared", Enabled: true},
|
|
{ChannelId: 562, Group: "auto-b", Model: "shared", Enabled: true},
|
|
{ChannelId: 562, Group: "auto-b", Model: "gpt-6-astra", Enabled: true},
|
|
{ChannelId: 563, Group: "auto-b", Model: "disabled-channel", Enabled: true},
|
|
{ChannelId: 562, Group: "auto-b", Model: "disabled-ability", Enabled: false},
|
|
{ChannelId: 564, Group: "private", Model: "unauthorized", Enabled: true},
|
|
}).Error)
|
|
for _, preference := range []string{"cost", "speed", "stable"} {
|
|
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
|
c.Request = httptest.NewRequest("GET", "/v1/models", nil)
|
|
common.SetContextKey(c, constant.ContextKeyUserGroup, "auto-a")
|
|
c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "smart", AutoPreference: preference})
|
|
c.Set("pool_failures", map[string]int{"auto-a": 2, "auto-b": 2})
|
|
require.Equal(t, []string{"gpt-6-astra", "model-a", "shared"}, AutomaticPoolModels(c))
|
|
}
|
|
require.NoError(t, db.Create(&model.Channel{Id: 565, Status: 1}).Error)
|
|
require.NoError(t, db.Create(&model.Ability{ChannelId: 565, Group: "new-group", Model: "new-model", Enabled: true}).Error)
|
|
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"auto-a":"","auto-b":"","new-group":""}`))
|
|
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"auto-a":1,"auto-b":1,"new-group":1}`))
|
|
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
|
c.Request = httptest.NewRequest("GET", "/v1/models", nil)
|
|
common.SetContextKey(c, constant.ContextKeyUserGroup, "auto-a")
|
|
require.Equal(t, []string{"gpt-6-astra", "model-a", "new-model", "shared"}, AutomaticPoolModels(c))
|
|
})
|
|
}
|
|
}
|