Du kannst nicht mehr als 25 Themen auswählen Themen müssen entweder mit einem Buchstaben oder einer Ziffer beginnen. Sie können Bindestriche („-“) enthalten und bis zu 35 Zeichen lang sein.
 
 
 

153 Zeilen
4.1 KiB

  1. package controller
  2. import (
  3. "bytes"
  4. "encoding/json"
  5. "net/http"
  6. "net/http/httptest"
  7. "testing"
  8. "github.com/QuantumNous/new-api/common"
  9. "github.com/QuantumNous/new-api/model"
  10. "github.com/glebarez/sqlite"
  11. "github.com/gin-gonic/gin"
  12. "github.com/stretchr/testify/assert"
  13. "github.com/stretchr/testify/require"
  14. "gorm.io/gorm"
  15. )
  16. func setupUserRateLimitControllerDB(t *testing.T) {
  17. t.Helper()
  18. db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
  19. require.NoError(t, err)
  20. sqlDB, _ := db.DB()
  21. sqlDB.SetMaxOpenConns(1)
  22. origDB := model.DB
  23. model.DB = db
  24. common.RedisEnabled = false
  25. require.NoError(t, db.AutoMigrate(&model.UserModelRateLimit{}))
  26. t.Cleanup(func() {
  27. model.DB = origDB
  28. sqlDB.Close()
  29. })
  30. }
  31. func setupUserRateLimitRouter() *gin.Engine {
  32. gin.SetMode(gin.TestMode)
  33. r := gin.New()
  34. g := r.Group("/api/user")
  35. {
  36. g.GET("/:id/rate-limits", GetUserRateLimits)
  37. g.PUT("/:id/rate-limits", SetUserRateLimits)
  38. }
  39. return r
  40. }
  41. func TestGetUserRateLimits_Empty(t *testing.T) {
  42. setupUserRateLimitControllerDB(t)
  43. router := setupUserRateLimitRouter()
  44. w := httptest.NewRecorder()
  45. req, _ := http.NewRequest(http.MethodGet, "/api/user/1/rate-limits", nil)
  46. router.ServeHTTP(w, req)
  47. assert.Equal(t, http.StatusOK, w.Code)
  48. var resp map[string]interface{}
  49. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
  50. assert.True(t, resp["success"].(bool))
  51. // data 应为空列表(非 nil)
  52. data, ok := resp["data"]
  53. require.True(t, ok, "response should have data key")
  54. assert.Empty(t, data)
  55. }
  56. func TestSetUserRateLimits_OK(t *testing.T) {
  57. setupUserRateLimitControllerDB(t)
  58. router := setupUserRateLimitRouter()
  59. items := []model.UserModelRateLimit{
  60. {Model: "gpt-4", Rpm: 60},
  61. }
  62. body, _ := json.Marshal(items)
  63. w := httptest.NewRecorder()
  64. req, _ := http.NewRequest(http.MethodPut, "/api/user/1/rate-limits", bytes.NewReader(body))
  65. req.Header.Set("Content-Type", "application/json")
  66. router.ServeHTTP(w, req)
  67. assert.Equal(t, http.StatusOK, w.Code)
  68. var resp map[string]interface{}
  69. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
  70. assert.True(t, resp["success"].(bool))
  71. // 验证数据确实写入
  72. list, err := model.GetUserModelRateLimits(1)
  73. require.NoError(t, err)
  74. require.Len(t, list, 1)
  75. assert.Equal(t, "gpt-4", list[0].Model)
  76. assert.Equal(t, 60, list[0].Rpm)
  77. }
  78. func TestSetUserRateLimits_NegativeRpm(t *testing.T) {
  79. setupUserRateLimitControllerDB(t)
  80. router := setupUserRateLimitRouter()
  81. items := []model.UserModelRateLimit{
  82. {Model: "gpt-4", Rpm: -1},
  83. }
  84. body, _ := json.Marshal(items)
  85. w := httptest.NewRecorder()
  86. req, _ := http.NewRequest(http.MethodPut, "/api/user/1/rate-limits", bytes.NewReader(body))
  87. req.Header.Set("Content-Type", "application/json")
  88. router.ServeHTTP(w, req)
  89. assert.Equal(t, http.StatusOK, w.Code)
  90. var resp map[string]interface{}
  91. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
  92. assert.False(t, resp["success"].(bool))
  93. }
  94. func TestGetUserRateLimits_InvalidId(t *testing.T) {
  95. setupUserRateLimitControllerDB(t)
  96. router := setupUserRateLimitRouter()
  97. w := httptest.NewRecorder()
  98. req, _ := http.NewRequest(http.MethodGet, "/api/user/notanid/rate-limits", nil)
  99. router.ServeHTTP(w, req)
  100. assert.Equal(t, http.StatusOK, w.Code)
  101. var resp map[string]interface{}
  102. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
  103. assert.False(t, resp["success"].(bool))
  104. }
  105. func TestSetUserRateLimits_EmptySlice(t *testing.T) {
  106. setupUserRateLimitControllerDB(t)
  107. router := setupUserRateLimitRouter()
  108. // 先写一条
  109. _ = model.SetUserModelRateLimits(2, []model.UserModelRateLimit{
  110. {UserId: 2, Model: "claude-3-opus", Rpm: 10},
  111. })
  112. // 用空 slice 清空
  113. body, _ := json.Marshal([]model.UserModelRateLimit{})
  114. w := httptest.NewRecorder()
  115. req, _ := http.NewRequest(http.MethodPut, "/api/user/2/rate-limits", bytes.NewReader(body))
  116. req.Header.Set("Content-Type", "application/json")
  117. router.ServeHTTP(w, req)
  118. assert.Equal(t, http.StatusOK, w.Code)
  119. var resp map[string]interface{}
  120. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
  121. assert.True(t, resp["success"].(bool))
  122. list, err := model.GetUserModelRateLimits(2)
  123. require.NoError(t, err)
  124. assert.Empty(t, list)
  125. }