package service import ( "fmt" "os" "strings" "testing" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/model" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "gorm.io/gorm" "gorm.io/gorm/logger" ) func setupDoubaoAssetChannelDB(t *testing.T) *gorm.DB { t.Helper() oldDB := model.DB oldSQLitePath := common.SQLitePath oldMemoryCacheEnabled := common.MemoryCacheEnabled oldIsMasterNode := common.IsMasterNode oldUsingSQLite := common.UsingSQLite oldUsingMySQL := common.UsingMySQL oldUsingPostgreSQL := common.UsingPostgreSQL oldSQLDSN, hadSQLDSN := os.LookupEnv("SQL_DSN") common.SQLitePath = "file:" + strings.NewReplacer("/", "_", " ", "_").Replace(t.Name()) + "?mode=memory&cache=shared" common.MemoryCacheEnabled = false common.IsMasterNode = false common.UsingSQLite = false common.UsingMySQL = false common.UsingPostgreSQL = false require.NoError(t, os.Setenv("SQL_DSN", "local")) require.NoError(t, model.InitDB()) model.DB = model.DB.Session(&gorm.Session{Logger: logger.Default.LogMode(logger.Silent)}) db := model.DB sqlDB, err := db.DB() require.NoError(t, err) sqlDB.SetMaxOpenConns(1) require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.Ability{}, &model.UserAssetChannel{})) t.Cleanup(func() { _ = sqlDB.Close() model.DB = oldDB common.SQLitePath = oldSQLitePath common.MemoryCacheEnabled = oldMemoryCacheEnabled common.IsMasterNode = oldIsMasterNode common.UsingSQLite = oldUsingSQLite common.UsingMySQL = oldUsingMySQL common.UsingPostgreSQL = oldUsingPostgreSQL if hadSQLDSN { _ = os.Setenv("SQL_DSN", oldSQLDSN) } else { _ = os.Unsetenv("SQL_DSN") } }) return db } func createDoubaoAssetChannelForTest(t *testing.T, db *gorm.DB, id int, channelType int, group string, key string, status int) { t.Helper() priority := int64(id) weight := uint(10) autoBan := 1 channel := &model.Channel{ Id: id, Type: channelType, Key: key, Status: status, Name: fmt.Sprintf("channel-%d", id), Group: group, Models: "seedance-2", Priority: &priority, Weight: &weight, AutoBan: &autoBan, CreatedTime: int64(id), } require.NoError(t, db.Create(channel).Error) } func createDoubaoAssetAbilityForTest(t *testing.T, db *gorm.DB, group string, modelName string, channelId int, enabled bool) { t.Helper() priority := int64(channelId) require.NoError(t, db.Create(&model.Ability{ Group: group, Model: modelName, ChannelId: channelId, Enabled: enabled, Priority: &priority, Weight: 10, }).Error) } func TestResolveDoubaoAssetChannel_PrefersBoundChannelWithKey(t *testing.T) { db := setupDoubaoAssetChannelDB(t) createDoubaoAssetChannelForTest(t, db, 1, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", "bound-key", common.ChannelStatusEnabled) createDoubaoAssetChannelForTest(t, db, 2, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", "auto-key", common.ChannelStatusEnabled) require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", 1)) channel, err := ResolveDoubaoAssetChannel(10, "default") require.NoError(t, err) require.NotNil(t, channel) assert.Equal(t, 1, channel.Id) assert.Equal(t, "bound-key", channel.Key) } func TestResolveDoubaoAssetChannel_AutoMatchesAndPersistsWhenUnbound(t *testing.T) { db := setupDoubaoAssetChannelDB(t) createDoubaoAssetChannelForTest(t, db, 1, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", "auto-key", common.ChannelStatusEnabled) channel, err := ResolveDoubaoAssetChannel(10, "default") require.NoError(t, err) require.NotNil(t, channel) assert.Equal(t, 1, channel.Id) binding, err := model.GetUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default") require.NoError(t, err) require.NotNil(t, binding) assert.Equal(t, 1, binding.ChannelId) } func TestResolveDoubaoAssetChannel_RebindsOpenAIMisbindToAiping(t *testing.T) { db := setupDoubaoAssetChannelDB(t) createDoubaoAssetChannelForTest(t, db, 1, constant.ChannelTypeOpenAI, "default", "openai-key", common.ChannelStatusEnabled) createDoubaoAssetChannelForTest(t, db, 2, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", "aiping-key", common.ChannelStatusEnabled) require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", 1)) channel, err := ResolveDoubaoAssetChannel(10, "default") require.NoError(t, err) require.NotNil(t, channel) assert.Equal(t, 2, channel.Id) binding, err := model.GetUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default") require.NoError(t, err) require.NotNil(t, binding) assert.Equal(t, 2, binding.ChannelId) } func TestResolveDoubaoAssetChannel_RebindsOnlyCurrentGroupWhenBoundGroupMismatches(t *testing.T) { db := setupDoubaoAssetChannelDB(t) createDoubaoAssetChannelForTest(t, db, 1, constant.ChannelTypeDoubaoVideoCompatibleAiping, "vip", "vip-key", common.ChannelStatusEnabled) createDoubaoAssetChannelForTest(t, db, 2, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", "default-key", common.ChannelStatusEnabled) require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", 1)) require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleAiping, "vip", 1)) channel, err := ResolveDoubaoAssetChannel(10, "default") require.NoError(t, err) require.NotNil(t, channel) assert.Equal(t, 2, channel.Id) defaultBinding, err := model.GetUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default") require.NoError(t, err) require.NotNil(t, defaultBinding) assert.Equal(t, 2, defaultBinding.ChannelId) vipBinding, err := model.GetUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleAiping, "vip") require.NoError(t, err) require.NotNil(t, vipBinding) assert.Equal(t, 1, vipBinding.ChannelId) } func TestResolveDoubaoAssetChannel_AutoMatchScansPastFirstPage(t *testing.T) { db := setupDoubaoAssetChannelDB(t) createDoubaoAssetChannelForTest(t, db, 1, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", "default-key", common.ChannelStatusEnabled) for i := 2; i <= DoubaoAssetChannelPageSize+1; i++ { createDoubaoAssetChannelForTest(t, db, i, constant.ChannelTypeDoubaoVideoCompatibleAiping, "vip", fmt.Sprintf("vip-key-%d", i), common.ChannelStatusEnabled) } channel, err := ResolveDoubaoAssetChannel(10, "default") require.NoError(t, err) require.NotNil(t, channel) assert.Equal(t, 1, channel.Id) } func TestResolveDoubaoAssetChannel_ReturnsErrorWhenBindingWriteFails(t *testing.T) { db := setupDoubaoAssetChannelDB(t) createDoubaoAssetChannelForTest(t, db, 1, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", "auto-key", common.ChannelStatusEnabled) const callbackName = "doubao_asset_test:fail_user_asset_channel_create" require.NoError(t, db.Callback().Create().Before("gorm:create").Register(callbackName, func(tx *gorm.DB) { if _, ok := tx.Statement.Dest.(*model.UserAssetChannel); ok { tx.AddError(assert.AnError) } })) channel, err := ResolveDoubaoAssetChannel(10, "default") require.Error(t, err) assert.Nil(t, channel) } func TestHasDoubaoVideoChannelForGroupModel(t *testing.T) { db := setupDoubaoAssetChannelDB(t) createDoubaoAssetChannelForTest(t, db, 1, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", "default-key", common.ChannelStatusEnabled) createDoubaoAssetChannelForTest(t, db, 2, constant.ChannelTypeDoubaoVideoCompatibleAiping, "vip", "vip-key", common.ChannelStatusEnabled) createDoubaoAssetChannelForTest(t, db, 3, constant.ChannelTypeOpenAI, "default", "openai-key", common.ChannelStatusEnabled) createDoubaoAssetAbilityForTest(t, db, "default", "seedance-2", 1, true) createDoubaoAssetAbilityForTest(t, db, "vip", "seedance-2", 2, true) createDoubaoAssetAbilityForTest(t, db, "default", "disabled-model", 1, false) createDoubaoAssetAbilityForTest(t, db, "default", "seedance-2", 3, true) assert.True(t, HasDoubaoVideoChannelForGroupModel("default", "seedance-2")) assert.False(t, HasDoubaoVideoChannelForGroupModel("default", "disabled-model")) assert.False(t, HasDoubaoVideoChannelForGroupModel("missing", "seedance-2")) } func TestAutoMatchDoubaoAssetChannel_IgnoresLegacyDoubaoVideoChannels(t *testing.T) { db := setupDoubaoAssetChannelDB(t) createDoubaoAssetChannelForTest(t, db, 1, constant.ChannelTypeDoubaoVideo, "default", "legacy-key", common.ChannelStatusEnabled) createDoubaoAssetChannelForTest(t, db, 2, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", "aiping-key", common.ChannelStatusEnabled) channel, err := AutoMatchDoubaoAssetChannel("default") require.NoError(t, err) require.NotNil(t, channel) assert.Equal(t, 2, channel.Id) assert.Equal(t, constant.ChannelTypeDoubaoVideoCompatibleAiping, channel.Type) }