package service import ( "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" ) func TestGetUserConcreteTokenGroups(t *testing.T) { db := setupDoubaoAssetChannelDB(t) for i, group := range []string{"vip", "default", "default", "auto", ""} { require.NoError(t, db.Create(&model.Token{Id: i + 1, UserId: 10, Key: "key-" + string(rune('0'+i)), Group: group}).Error) } groups, err := model.GetUserConcreteTokenGroups(10) require.NoError(t, err) assert.Equal(t, []string{"default", "vip"}, groups) } func TestGetVideoAssetChannelCandidatesFiltersGroupAndStatus(t *testing.T) { db := setupDoubaoAssetChannelDB(t) createDoubaoAssetChannelForTest(t, db, 7, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", "key", common.ChannelStatusEnabled) createDoubaoAssetChannelForTest(t, db, 16, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "default", "key", common.ChannelStatusEnabled) createDoubaoAssetChannelForTest(t, db, 17, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "default", "key", common.ChannelStatusAutoDisabled) createDoubaoAssetChannelForTest(t, db, 18, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "vip", "key", common.ChannelStatusEnabled) createDoubaoAssetChannelForTest(t, db, 19, constant.ChannelTypeChinaMobileSeedance, "default", "\n", common.ChannelStatusEnabled) candidates, err := GetVideoAssetChannelCandidates("default", VideoAssetFamilySeedance) require.NoError(t, err) require.Len(t, candidates, 2) assert.Equal(t, []int{7, 16}, []int{candidates[0].Id, candidates[1].Id}) } func TestClearVideoAssetChannelBindingClearsOnlyFamily(t *testing.T) { db := setupDoubaoAssetChannelDB(t) createDoubaoAssetChannelForTest(t, db, 7, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", "key", common.ChannelStatusEnabled) createDoubaoAssetChannelForTest(t, db, 59, constant.ChannelTypeKlingAiping, "default", "key", common.ChannelStatusEnabled) require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", 7)) require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeKlingAiping, "default", 59)) require.NoError(t, ClearVideoAssetChannelBinding(10, "default", VideoAssetFamilySeedance)) seedanceBindings, err := model.GetUserAssetChannelsByTypes(10, VideoAssetChannelTypesForFamily(VideoAssetFamilySeedance), "default") require.NoError(t, err) assert.Empty(t, seedanceBindings) klingBinding, err := model.GetUserAssetChannel(10, constant.ChannelTypeKlingAiping, "default") require.NoError(t, err) assert.NotNil(t, klingBinding) }