new-api/service/channel_select_test.go

package service

import (
	"net/http"
	"net/http/httptest"
	"os"
	"strings"
	"testing"

	"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/gin-gonic/gin"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
)

func buildRetryContext(t *testing.T) *gin.Context {
	t.Helper()
	gin.SetMode(gin.TestMode)
	w := httptest.NewRecorder()
	c, _ := gin.CreateTestContext(w)
	req, err := http.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
	require.NoError(t, err)
	c.Request = req
	return c
}

func setupServiceChannelSelectDB(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{}))

	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 createServiceChannelSelectChannel(t *testing.T, db *gorm.DB, id int, channelType int, group string, modelName string) {
	t.Helper()

	priority := int64(10)
	weight := uint(10)
	autoBan := 1
	require.NoError(t, db.Create(&model.Channel{
		Id:       id,
		Type:     channelType,
		Key:      "key",
		Status:   common.ChannelStatusEnabled,
		Name:     "channel",
		Group:    group,
		Models:   modelName,
		Priority: &priority,
		Weight:   &weight,
		AutoBan:  &autoBan,
	}).Error)
	require.NoError(t, db.Create(&model.Ability{
		Group:     group,
		Model:     modelName,
		ChannelId: id,
		Enabled:   true,
		Priority:  &priority,
		Weight:    10,
	}).Error)
}

// --- RetryParam tests ---

func TestRetryParam_IncreaseRetry(t *testing.T) {
	p := &RetryParam{}
	assert.Equal(t, 0, p.GetRetry())

	p.IncreaseRetry()
	assert.Equal(t, 1, p.GetRetry())

	p.IncreaseRetry()
	assert.Equal(t, 2, p.GetRetry())
}

func TestRetryParam_ResetRetryNextTry(t *testing.T) {
	p := &RetryParam{}
	p.IncreaseRetry()
	assert.Equal(t, 1, p.GetRetry())

	// ResetRetryNextTry should cause next IncreaseRetry to be a no-op
	p.ResetRetryNextTry()
	p.IncreaseRetry()
	assert.Equal(t, 1, p.GetRetry()) // still 1

	p.IncreaseRetry()
	assert.Equal(t, 2, p.GetRetry()) // now increases
}

func TestRetryParam_SetRetry(t *testing.T) {
	p := &RetryParam{}
	p.SetRetry(5)
	assert.Equal(t, 5, p.GetRetry())
}

// --- Auto group selection (no DB, tests with no channels available) ---

func TestCacheGetRandomSatisfiedChannel_AutoGroup_NoAutoGroups(t *testing.T) {
	require.NoError(t, setting.UpdateAutoGroupsByJsonString(`[]`))
	t.Cleanup(func() {
		_ = setting.UpdateAutoGroupsByJsonString(`["default"]`)
		_ = setting.UpdateUserUsableGroupsByJSONString(`{"default":"默认分组","vip":"vip分组"}`)
	})

	c := buildRetryContext(t)
	p := &RetryParam{
		Ctx:        c,
		TokenGroup: "auto",
		ModelName:  "gpt-4",
	}

	ch, group, err := CacheGetRandomSatisfiedChannel(p)
	assert.Nil(t, ch)
	assert.Error(t, err)
	assert.Equal(t, "auto", group)
}

func TestCacheGetRandomSatisfiedChannelFiltersAllowedChannelTypes(t *testing.T) {
	db := setupServiceChannelSelectDB(t)
	createServiceChannelSelectChannel(t, db, 101, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", "Doubao-Seedance-2.0")
	createServiceChannelSelectChannel(t, db, 102, constant.ChannelTypeOpenAI, "default", "Doubao-Seedance-2.0")

	c := buildRetryContext(t)
	channel, group, err := CacheGetRandomSatisfiedChannel(&RetryParam{
		Ctx:                 c,
		TokenGroup:          "default",
		ModelName:           "Doubao-Seedance-2.0",
		Retry:               common.GetPointer(0),
		AllowedChannelTypes: []int{constant.ChannelTypeDoubaoVideoCompatibleAiping, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun},
	})

	require.NoError(t, err)
	require.NotNil(t, channel)
	assert.Equal(t, "default", group)
	assert.Equal(t, 101, channel.Id)
}

// --- GetUserAutoGroup tests ---

func TestGetUserAutoGroup_Basic(t *testing.T) {
	require.NoError(t, setting.UpdateAutoGroupsByJsonString(`["default","vip"]`))
	require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"默认分组","vip":"vip分组"}`))
	t.Cleanup(func() {
		_ = setting.UpdateAutoGroupsByJsonString(`["default"]`)
		_ = setting.UpdateUserUsableGroupsByJSONString(`{"default":"默认分组","vip":"vip分组"}`)
	})

	// default user can use both groups
	autoGroups := GetUserAutoGroup("default")
	assert.Equal(t, []string{"default", "vip"}, autoGroups)
}

func TestGetUserAutoGroup_VipUserWithExtraGroups(t *testing.T) {
	require.NoError(t, setting.UpdateAutoGroupsByJsonString(`["default","vip","premium"]`))
	require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"默认分组","vip":"vip分组","premium":"高级分组"}`))
	t.Cleanup(func() {
		_ = setting.UpdateAutoGroupsByJsonString(`["default"]`)
		_ = setting.UpdateUserUsableGroupsByJSONString(`{"default":"默认分组","vip":"vip分组"}`)
	})

	autoGroups := GetUserAutoGroup("vip")
	assert.Equal(t, []string{"default", "vip", "premium"}, autoGroups)
}

// --- GetUserGroupRatio tests ---

func TestGetUserGroupRatio_SpecialRatio(t *testing.T) {
	// When user has a special ratio for a group, it should be used
	// The default groupGroupRatio has vip->edit_this = 0.9
	ratio := GetUserGroupRatio("vip", "edit_this")
	assert.Equal(t, 0.9, ratio)
}

func TestGetUserGroupRatio_FallbackToGroupRatio(t *testing.T) {
	// When no special ratio exists, fall back to group ratio
	ratio := GetUserGroupRatio("default", "default")
	assert.Equal(t, 1.0, ratio) // default group has ratio 1
}