Non puoi selezionare più di 25 argomenti Gli argomenti devono iniziare con una lettera o un numero, possono includere trattini ('-') e possono essere lunghi fino a 35 caratteri.
 
 
 

121 righe
3.5 KiB

  1. package service
  2. import (
  3. "net/http"
  4. "net/http/httptest"
  5. "testing"
  6. "github.com/QuantumNous/new-api/setting"
  7. "github.com/gin-gonic/gin"
  8. "github.com/stretchr/testify/assert"
  9. "github.com/stretchr/testify/require"
  10. )
  11. func buildRetryContext(t *testing.T) *gin.Context {
  12. t.Helper()
  13. gin.SetMode(gin.TestMode)
  14. w := httptest.NewRecorder()
  15. c, _ := gin.CreateTestContext(w)
  16. req, err := http.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
  17. require.NoError(t, err)
  18. c.Request = req
  19. return c
  20. }
  21. // --- RetryParam tests ---
  22. func TestRetryParam_IncreaseRetry(t *testing.T) {
  23. p := &RetryParam{}
  24. assert.Equal(t, 0, p.GetRetry())
  25. p.IncreaseRetry()
  26. assert.Equal(t, 1, p.GetRetry())
  27. p.IncreaseRetry()
  28. assert.Equal(t, 2, p.GetRetry())
  29. }
  30. func TestRetryParam_ResetRetryNextTry(t *testing.T) {
  31. p := &RetryParam{}
  32. p.IncreaseRetry()
  33. assert.Equal(t, 1, p.GetRetry())
  34. // ResetRetryNextTry should cause next IncreaseRetry to be a no-op
  35. p.ResetRetryNextTry()
  36. p.IncreaseRetry()
  37. assert.Equal(t, 1, p.GetRetry()) // still 1
  38. p.IncreaseRetry()
  39. assert.Equal(t, 2, p.GetRetry()) // now increases
  40. }
  41. func TestRetryParam_SetRetry(t *testing.T) {
  42. p := &RetryParam{}
  43. p.SetRetry(5)
  44. assert.Equal(t, 5, p.GetRetry())
  45. }
  46. // --- Auto group selection (no DB, tests with no channels available) ---
  47. func TestCacheGetRandomSatisfiedChannel_AutoGroup_NoAutoGroups(t *testing.T) {
  48. require.NoError(t, setting.UpdateAutoGroupsByJsonString(`[]`))
  49. t.Cleanup(func() {
  50. _ = setting.UpdateAutoGroupsByJsonString(`["default"]`)
  51. _ = setting.UpdateUserUsableGroupsByJSONString(`{"default":"默认分组","vip":"vip分组"}`)
  52. })
  53. c := buildRetryContext(t)
  54. p := &RetryParam{
  55. Ctx: c,
  56. TokenGroup: "auto",
  57. ModelName: "gpt-4",
  58. }
  59. ch, group, err := CacheGetRandomSatisfiedChannel(p)
  60. assert.Nil(t, ch)
  61. assert.Error(t, err)
  62. assert.Equal(t, "auto", group)
  63. }
  64. // --- GetUserAutoGroup tests ---
  65. func TestGetUserAutoGroup_Basic(t *testing.T) {
  66. require.NoError(t, setting.UpdateAutoGroupsByJsonString(`["default","vip"]`))
  67. require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"默认分组","vip":"vip分组"}`))
  68. t.Cleanup(func() {
  69. _ = setting.UpdateAutoGroupsByJsonString(`["default"]`)
  70. _ = setting.UpdateUserUsableGroupsByJSONString(`{"default":"默认分组","vip":"vip分组"}`)
  71. })
  72. // default user can use both groups
  73. autoGroups := GetUserAutoGroup("default")
  74. assert.Equal(t, []string{"default", "vip"}, autoGroups)
  75. }
  76. func TestGetUserAutoGroup_VipUserWithExtraGroups(t *testing.T) {
  77. require.NoError(t, setting.UpdateAutoGroupsByJsonString(`["default","vip","premium"]`))
  78. require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"默认分组","vip":"vip分组","premium":"高级分组"}`))
  79. t.Cleanup(func() {
  80. _ = setting.UpdateAutoGroupsByJsonString(`["default"]`)
  81. _ = setting.UpdateUserUsableGroupsByJSONString(`{"default":"默认分组","vip":"vip分组"}`)
  82. })
  83. autoGroups := GetUserAutoGroup("vip")
  84. assert.Equal(t, []string{"default", "vip", "premium"}, autoGroups)
  85. }
  86. // --- GetUserGroupRatio tests ---
  87. func TestGetUserGroupRatio_SpecialRatio(t *testing.T) {
  88. // When user has a special ratio for a group, it should be used
  89. // The default groupGroupRatio has vip->edit_this = 0.9
  90. ratio := GetUserGroupRatio("vip", "edit_this")
  91. assert.Equal(t, 0.9, ratio)
  92. }
  93. func TestGetUserGroupRatio_FallbackToGroupRatio(t *testing.T) {
  94. // When no special ratio exists, fall back to group ratio
  95. ratio := GetUserGroupRatio("default", "default")
  96. assert.Equal(t, 1.0, ratio) // default group has ratio 1
  97. }