package service import ( "net/http" "net/http/httptest" "testing" "github.com/QuantumNous/new-api/setting" "github.com/gin-gonic/gin" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) 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 } // --- 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) } // --- 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 }