|
- 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
- }
|