You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

217 lines
5.8 KiB

  1. package middleware
  2. import (
  3. "net/http"
  4. "net/http/httptest"
  5. "sync"
  6. "sync/atomic"
  7. "testing"
  8. "github.com/QuantumNous/new-api/common"
  9. "github.com/QuantumNous/new-api/model"
  10. "github.com/QuantumNous/new-api/setting"
  11. "github.com/gin-gonic/gin"
  12. "github.com/glebarez/sqlite"
  13. "github.com/stretchr/testify/assert"
  14. "github.com/stretchr/testify/require"
  15. "gorm.io/gorm"
  16. )
  17. func setupTestDB(t *testing.T) {
  18. t.Helper()
  19. db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
  20. require.NoError(t, err)
  21. sqlDB, _ := db.DB()
  22. sqlDB.SetMaxOpenConns(1)
  23. origDB := model.DB
  24. model.DB = db
  25. require.NoError(t, db.AutoMigrate(&model.UserModelRateLimit{}))
  26. t.Cleanup(func() {
  27. model.DB = origDB
  28. sqlDB.Close()
  29. })
  30. }
  31. // buildTestRouter 创建用于测试的 gin Engine,前置中间件注入 userId 和 modelName
  32. func buildTestRouter(userId int, modelName string) *gin.Engine {
  33. gin.SetMode(gin.TestMode)
  34. r := gin.New()
  35. r.POST("/test",
  36. func(c *gin.Context) {
  37. c.Set("id", userId)
  38. c.Set("original_model", modelName)
  39. c.Next()
  40. },
  41. UserModelRateLimit(),
  42. func(c *gin.Context) {
  43. c.JSON(http.StatusOK, gin.H{"ok": true})
  44. },
  45. )
  46. return r
  47. }
  48. func doRequest(r *gin.Engine) *httptest.ResponseRecorder {
  49. w := httptest.NewRecorder()
  50. req := httptest.NewRequest(http.MethodPost, "/test", nil)
  51. r.ServeHTTP(w, req)
  52. return w
  53. }
  54. // TestUserModelRateLimit_Disabled:全局开关关闭时即使有限流配置也应放行
  55. func TestUserModelRateLimit_Disabled(t *testing.T) {
  56. setupTestDB(t)
  57. common.RedisEnabled = false
  58. setting.ModelRequestRateLimitEnabled = false
  59. // 插入 rpm=1 的严格限制
  60. require.NoError(t, model.DB.Create(&model.UserModelRateLimit{
  61. UserId: 42, Model: "gpt-4", Rpm: 1,
  62. }).Error)
  63. r := buildTestRouter(42, "gpt-4")
  64. // 连续发两次请求,都应该 200
  65. w1 := doRequest(r)
  66. assert.Equal(t, http.StatusOK, w1.Code)
  67. w2 := doRequest(r)
  68. assert.Equal(t, http.StatusOK, w2.Code)
  69. }
  70. // TestUserModelRateLimit_NoModel:original_model 为空时放行
  71. func TestUserModelRateLimit_NoModel(t *testing.T) {
  72. setupTestDB(t)
  73. common.RedisEnabled = false
  74. setting.ModelRequestRateLimitEnabled = true
  75. r := buildTestRouter(42, "")
  76. w := doRequest(r)
  77. assert.Equal(t, http.StatusOK, w.Code)
  78. }
  79. // TestUserModelRateLimit_NoConfig:DB 中无该用户+模型的配置时放行
  80. func TestUserModelRateLimit_NoConfig(t *testing.T) {
  81. setupTestDB(t)
  82. common.RedisEnabled = false
  83. setting.ModelRequestRateLimitEnabled = true
  84. setting.ModelRequestRateLimitDurationMinutes = 1
  85. // DB 中没有任何配置
  86. r := buildTestRouter(99, "gpt-4")
  87. w := doRequest(r)
  88. assert.Equal(t, http.StatusOK, w.Code)
  89. }
  90. // TestUserModelRateLimit_UnderLimit:rpm=2,发 1 次请求,放行
  91. func TestUserModelRateLimit_UnderLimit(t *testing.T) {
  92. setupTestDB(t)
  93. common.RedisEnabled = false
  94. setting.ModelRequestRateLimitEnabled = true
  95. setting.ModelRequestRateLimitDurationMinutes = 1
  96. require.NoError(t, model.DB.Create(&model.UserModelRateLimit{
  97. UserId: 10, Model: "claude-3", Rpm: 2,
  98. }).Error)
  99. r := buildTestRouter(10, "claude-3")
  100. w := doRequest(r)
  101. assert.Equal(t, http.StatusOK, w.Code)
  102. }
  103. // TestUserModelRateLimit_OverLimit:rpm=2,连续发 3 次请求,第 3 次被 429
  104. func TestUserModelRateLimit_OverLimit(t *testing.T) {
  105. setupTestDB(t)
  106. common.RedisEnabled = false
  107. setting.ModelRequestRateLimitEnabled = true
  108. setting.ModelRequestRateLimitDurationMinutes = 1
  109. require.NoError(t, model.DB.Create(&model.UserModelRateLimit{
  110. UserId: 20, Model: "gpt-4o", Rpm: 2,
  111. }).Error)
  112. r := buildTestRouter(20, "gpt-4o")
  113. w1 := doRequest(r)
  114. assert.Equal(t, http.StatusOK, w1.Code, "第 1 次请求应放行")
  115. w2 := doRequest(r)
  116. assert.Equal(t, http.StatusOK, w2.Code, "第 2 次请求应放行")
  117. w3 := doRequest(r)
  118. assert.Equal(t, http.StatusTooManyRequests, w3.Code, "第 3 次请求应被限流")
  119. }
  120. // TestUserModelRateLimit_ConcurrentMemory:内存路径并发安全验证。
  121. // RPM=2,并发 20 个请求,预期最多 2 个 200,其余 429。
  122. func TestUserModelRateLimit_ConcurrentMemory(t *testing.T) {
  123. setupTestDB(t)
  124. common.RedisEnabled = false
  125. setting.ModelRequestRateLimitEnabled = true
  126. setting.ModelRequestRateLimitDurationMinutes = 1
  127. require.NoError(t, model.DB.Create(&model.UserModelRateLimit{
  128. UserId: 30, Model: "deepseek-r1", Rpm: 2,
  129. }).Error)
  130. r := buildTestRouter(30, "deepseek-r1")
  131. var successCount int64
  132. var wg sync.WaitGroup
  133. totalRequests := 20
  134. for i := 0; i < totalRequests; i++ {
  135. wg.Add(1)
  136. go func() {
  137. defer wg.Done()
  138. w := doRequest(r)
  139. if w.Code == http.StatusOK {
  140. atomic.AddInt64(&successCount, 1)
  141. }
  142. }()
  143. }
  144. wg.Wait()
  145. sc := int(atomic.LoadInt64(&successCount))
  146. assert.LessOrEqual(t, sc, 2, "并发请求中最多 %d 个应成功,实际 %d 个", 2, sc)
  147. assert.GreaterOrEqual(t, sc, 1, "至少应有 1 个请求成功")
  148. }
  149. // TestUserModelRateLimit_UserIsolation:不同用户之间的限流互不影响
  150. func TestUserModelRateLimit_UserIsolation(t *testing.T) {
  151. setupTestDB(t)
  152. common.RedisEnabled = false
  153. setting.ModelRequestRateLimitEnabled = true
  154. setting.ModelRequestRateLimitDurationMinutes = 1
  155. // 两个用户都设置 RPM=2
  156. require.NoError(t, model.DB.Create(&model.UserModelRateLimit{
  157. UserId: 100, Model: "gpt-4", Rpm: 2,
  158. }).Error)
  159. require.NoError(t, model.DB.Create(&model.UserModelRateLimit{
  160. UserId: 101, Model: "gpt-4", Rpm: 2,
  161. }).Error)
  162. r1 := buildTestRouter(100, "gpt-4")
  163. r2 := buildTestRouter(101, "gpt-4")
  164. // 用户 100 用完配额
  165. w1 := doRequest(r1)
  166. assert.Equal(t, http.StatusOK, w1.Code)
  167. w2 := doRequest(r1)
  168. assert.Equal(t, http.StatusOK, w2.Code)
  169. w3 := doRequest(r1)
  170. assert.Equal(t, http.StatusTooManyRequests, w3.Code)
  171. // 用户 101 应该不受影响
  172. w4 := doRequest(r2)
  173. assert.Equal(t, http.StatusOK, w4.Code)
  174. w5 := doRequest(r2)
  175. assert.Equal(t, http.StatusOK, w5.Code)
  176. w6 := doRequest(r2)
  177. assert.Equal(t, http.StatusTooManyRequests, w6.Code)
  178. }