feat: sync deployed v1.0.0-rc.38-custom.67

2026-10-06 00:07:10 +08:00
1 parent d2e03ab7e8
commit ffc8cf56db
@@ -0,0 +1,97 @@
+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))
+		})
+	}
+}