new-api/model/channel_select_test.go

package model

import (
	"testing"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/constant"
	"github.com/glebarez/sqlite"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"gorm.io/gorm"
)

func setupChannelSelectDB(t *testing.T) *gorm.DB {
	t.Helper()

	db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
	require.NoError(t, err)

	sqlDB, err := db.DB()
	require.NoError(t, err)
	sqlDB.SetMaxOpenConns(1)

	origDB := DB
	origUsingSQLite := common.UsingSQLite
	origRedisEnabled := common.RedisEnabled
	origMemoryCache := common.MemoryCacheEnabled
	origGroupCol := commonGroupCol
	origTrueVal := commonTrueVal

	DB = db
	common.UsingSQLite = true
	common.RedisEnabled = false
	common.MemoryCacheEnabled = false
	// Set column quoting for SQLite
	commonGroupCol = "`group`"
	commonTrueVal = "1"

	require.NoError(t, db.AutoMigrate(&Channel{}, &Ability{}))

	t.Cleanup(func() {
		DB = origDB
		common.UsingSQLite = origUsingSQLite
		common.RedisEnabled = origRedisEnabled
		common.MemoryCacheEnabled = origMemoryCache
		commonGroupCol = origGroupCol
		commonTrueVal = origTrueVal
		sqlDB.Close()
	})

	return db
}

func int64Ptr(v int64) *int64 { return &v }
func intPtr(v int) *int       { return &v }
func uintPtr(v uint) *uint    { return &v }

func createChannelForTest(t *testing.T, db *gorm.DB, id int, name, group, models string, priority int64) {
	t.Helper()
	ch := Channel{
		Id:          id,
		Type:        1,
		Key:         "test-key-" + name,
		Status:      common.ChannelStatusEnabled,
		Name:        name,
		Group:       group,
		Models:      models,
		Priority:    int64Ptr(priority),
		Weight:      uintPtr(10),
		AutoBan:     intPtr(1),
		CreatedTime: 1,
	}
	require.NoError(t, db.Create(&ch).Error)
}

func createChannelWithTypeForTest(t *testing.T, db *gorm.DB, id int, name, group, models string, channelType int, priority int64) {
	t.Helper()
	ch := Channel{
		Id:          id,
		Type:        channelType,
		Key:         "test-key-" + name,
		Status:      common.ChannelStatusEnabled,
		Name:        name,
		Group:       group,
		Models:      models,
		Priority:    int64Ptr(priority),
		Weight:      uintPtr(10),
		AutoBan:     intPtr(1),
		CreatedTime: 1,
	}
	require.NoError(t, db.Create(&ch).Error)
}

func createAbilityForTest(t *testing.T, db *gorm.DB, group, modelName string, channelId int, priority int64, enabled bool) {
	t.Helper()
	ability := Ability{
		Group:     group,
		Model:     modelName,
		ChannelId: channelId,
		Enabled:   enabled,
		Priority:  int64Ptr(priority),
		Weight:    10,
	}
	require.NoError(t, db.Create(&ability).Error)
}

// --- GetChannel (DB path) tests ---

func TestGetChannel_BasicSelection(t *testing.T) {
	db := setupChannelSelectDB(t)

	createChannelForTest(t, db, 1, "ch1", "default", "gpt-4", 10)
	createAbilityForTest(t, db, "default", "gpt-4", 1, 10, true)

	ch, err := GetChannel("default", "gpt-4", 0)
	require.NoError(t, err)
	assert.NotNil(t, ch)
	assert.Equal(t, 1, ch.Id)
}

func TestGetChannel_NoChannelForModel(t *testing.T) {
	db := setupChannelSelectDB(t)

	createChannelForTest(t, db, 1, "ch1", "default", "gpt-3.5", 10)
	createAbilityForTest(t, db, "default", "gpt-3.5", 1, 10, true)

	ch, err := GetChannel("default", "gpt-4", 0)
	assert.NoError(t, err)
	assert.Nil(t, ch)
}

func TestGetChannel_PrioritySelection(t *testing.T) {
	db := setupChannelSelectDB(t)

	// Two channels with different priorities
	createChannelForTest(t, db, 1, "ch-low", "default", "gpt-4", 0)
	createChannelForTest(t, db, 2, "ch-high", "default", "gpt-4", 10)
	createAbilityForTest(t, db, "default", "gpt-4", 1, 0, true)
	createAbilityForTest(t, db, "default", "gpt-4", 2, 10, true)

	// Retry=0 should get highest priority (10)
	ch, err := GetChannel("default", "gpt-4", 0)
	require.NoError(t, err)
	assert.NotNil(t, ch)
	assert.Equal(t, 2, ch.Id) // ch-high has priority 10

	// Retry=1 should get next priority (0)
	ch, err = GetChannel("default", "gpt-4", 1)
	require.NoError(t, err)
	assert.NotNil(t, ch)
	assert.Equal(t, 1, ch.Id) // ch-low has priority 0
}

func TestGetChannel_DifferentGroups(t *testing.T) {
	db := setupChannelSelectDB(t)

	// Channel in "default" group
	createChannelForTest(t, db, 1, "ch-default", "default", "gpt-4", 10)
	createAbilityForTest(t, db, "default", "gpt-4", 1, 10, true)

	// Channel in "vip" group
	createChannelForTest(t, db, 2, "ch-vip", "vip", "gpt-4", 10)
	createAbilityForTest(t, db, "vip", "gpt-4", 2, 10, true)

	// Query default group
	ch, err := GetChannel("default", "gpt-4", 0)
	require.NoError(t, err)
	assert.Equal(t, 1, ch.Id)

	// Query vip group
	ch, err = GetChannel("vip", "gpt-4", 0)
	require.NoError(t, err)
	assert.Equal(t, 2, ch.Id)
}

func TestGetChannel_DisabledAbility(t *testing.T) {
	db := setupChannelSelectDB(t)

	createChannelForTest(t, db, 1, "ch1", "default", "gpt-4", 10)
	createAbilityForTest(t, db, "default", "gpt-4", 1, 10, false) // disabled

	ch, err := GetChannel("default", "gpt-4", 0)
	assert.NoError(t, err)
	assert.Nil(t, ch) // no enabled ability
}

func TestGetChannel_MultipleChannelsWeighted(t *testing.T) {
	db := setupChannelSelectDB(t)

	createChannelForTest(t, db, 1, "ch1", "default", "gpt-4", 10)
	createChannelForTest(t, db, 2, "ch2", "default", "gpt-4", 10)
	createAbilityForTest(t, db, "default", "gpt-4", 1, 10, true)
	createAbilityForTest(t, db, "default", "gpt-4", 2, 10, true)

	// Run multiple times to verify both channels can be selected
	found := make(map[int]bool)
	for i := 0; i < 50; i++ {
		ch, err := GetChannel("default", "gpt-4", 0)
		require.NoError(t, err)
		require.NotNil(t, ch)
		found[ch.Id] = true
	}
	// With equal weights, both should appear at least once in 50 tries
	assert.True(t, found[1], "channel 1 should be selected at least once")
	assert.True(t, found[2], "channel 2 should be selected at least once")
}

func TestGetRandomSatisfiedChannelWithFilter_AllCandidatesRejected(t *testing.T) {
	db := setupChannelSelectDB(t)

	createChannelForTest(t, db, 1, "ch1", "default", "seedance-2", 10)
	createChannelForTest(t, db, 2, "ch2", "default", "seedance-2", 10)
	createAbilityForTest(t, db, "default", "seedance-2", 1, 10, true)
	createAbilityForTest(t, db, "default", "seedance-2", 2, 10, true)

	ch, err := GetRandomSatisfiedChannelWithFilter("default", "seedance-2", 0, func(*Channel) (bool, error) {
		return false, nil
	})
	require.NoError(t, err)
	assert.Nil(t, ch)
}

func TestGetRandomSatisfiedChannelWithFilter_SelectsOnlyAcceptedCandidate(t *testing.T) {
	db := setupChannelSelectDB(t)

	createChannelWithTypeForTest(t, db, 1, "openai", "default", "seedance-2", constant.ChannelTypeOpenAI, 10)
	createChannelWithTypeForTest(t, db, 2, "doubao", "default", "seedance-2", constant.ChannelTypeDoubaoVideo, 10)
	createAbilityForTest(t, db, "default", "seedance-2", 1, 10, true)
	createAbilityForTest(t, db, "default", "seedance-2", 2, 10, true)

	for i := 0; i < 20; i++ {
		ch, err := GetRandomSatisfiedChannelWithFilter("default", "seedance-2", 0, func(channel *Channel) (bool, error) {
			return channel.Type == constant.ChannelTypeDoubaoVideo, nil
		})
		require.NoError(t, err)
		require.NotNil(t, ch)
		assert.Equal(t, 2, ch.Id)
	}
}