|
- package middleware
-
- import (
- "net/http"
- "net/http/httptest"
- "sync"
- "sync/atomic"
- "testing"
-
- "github.com/QuantumNous/new-api/common"
- "github.com/QuantumNous/new-api/model"
- "github.com/QuantumNous/new-api/setting"
- "github.com/gin-gonic/gin"
- "github.com/glebarez/sqlite"
- "github.com/stretchr/testify/assert"
- "github.com/stretchr/testify/require"
- "gorm.io/gorm"
- )
-
- func setupTestDB(t *testing.T) {
- t.Helper()
- db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
- require.NoError(t, err)
- sqlDB, _ := db.DB()
- sqlDB.SetMaxOpenConns(1)
-
- origDB := model.DB
- model.DB = db
- require.NoError(t, db.AutoMigrate(&model.UserModelRateLimit{}))
-
- t.Cleanup(func() {
- model.DB = origDB
- sqlDB.Close()
- })
- }
-
- // buildTestRouter 创建用于测试的 gin Engine,前置中间件注入 userId 和 modelName
- func buildTestRouter(userId int, modelName string) *gin.Engine {
- gin.SetMode(gin.TestMode)
- r := gin.New()
- r.POST("/test",
- func(c *gin.Context) {
- c.Set("id", userId)
- c.Set("original_model", modelName)
- c.Next()
- },
- UserModelRateLimit(),
- func(c *gin.Context) {
- c.JSON(http.StatusOK, gin.H{"ok": true})
- },
- )
- return r
- }
-
- func doRequest(r *gin.Engine) *httptest.ResponseRecorder {
- w := httptest.NewRecorder()
- req := httptest.NewRequest(http.MethodPost, "/test", nil)
- r.ServeHTTP(w, req)
- return w
- }
-
- // TestUserModelRateLimit_Disabled:全局开关关闭时即使有限流配置也应放行
- func TestUserModelRateLimit_Disabled(t *testing.T) {
- setupTestDB(t)
- common.RedisEnabled = false
- setting.ModelRequestRateLimitEnabled = false
-
- // 插入 rpm=1 的严格限制
- require.NoError(t, model.DB.Create(&model.UserModelRateLimit{
- UserId: 42, Model: "gpt-4", Rpm: 1,
- }).Error)
-
- r := buildTestRouter(42, "gpt-4")
-
- // 连续发两次请求,都应该 200
- w1 := doRequest(r)
- assert.Equal(t, http.StatusOK, w1.Code)
- w2 := doRequest(r)
- assert.Equal(t, http.StatusOK, w2.Code)
- }
-
- // TestUserModelRateLimit_NoModel:original_model 为空时放行
- func TestUserModelRateLimit_NoModel(t *testing.T) {
- setupTestDB(t)
- common.RedisEnabled = false
- setting.ModelRequestRateLimitEnabled = true
-
- r := buildTestRouter(42, "")
-
- w := doRequest(r)
- assert.Equal(t, http.StatusOK, w.Code)
- }
-
- // TestUserModelRateLimit_NoConfig:DB 中无该用户+模型的配置时放行
- func TestUserModelRateLimit_NoConfig(t *testing.T) {
- setupTestDB(t)
- common.RedisEnabled = false
- setting.ModelRequestRateLimitEnabled = true
- setting.ModelRequestRateLimitDurationMinutes = 1
-
- // DB 中没有任何配置
- r := buildTestRouter(99, "gpt-4")
-
- w := doRequest(r)
- assert.Equal(t, http.StatusOK, w.Code)
- }
-
- // TestUserModelRateLimit_UnderLimit:rpm=2,发 1 次请求,放行
- func TestUserModelRateLimit_UnderLimit(t *testing.T) {
- setupTestDB(t)
- common.RedisEnabled = false
- setting.ModelRequestRateLimitEnabled = true
- setting.ModelRequestRateLimitDurationMinutes = 1
-
- require.NoError(t, model.DB.Create(&model.UserModelRateLimit{
- UserId: 10, Model: "claude-3", Rpm: 2,
- }).Error)
-
- r := buildTestRouter(10, "claude-3")
-
- w := doRequest(r)
- assert.Equal(t, http.StatusOK, w.Code)
- }
-
- // TestUserModelRateLimit_OverLimit:rpm=2,连续发 3 次请求,第 3 次被 429
- func TestUserModelRateLimit_OverLimit(t *testing.T) {
- setupTestDB(t)
- common.RedisEnabled = false
- setting.ModelRequestRateLimitEnabled = true
- setting.ModelRequestRateLimitDurationMinutes = 1
-
- require.NoError(t, model.DB.Create(&model.UserModelRateLimit{
- UserId: 20, Model: "gpt-4o", Rpm: 2,
- }).Error)
-
- r := buildTestRouter(20, "gpt-4o")
-
- w1 := doRequest(r)
- assert.Equal(t, http.StatusOK, w1.Code, "第 1 次请求应放行")
-
- w2 := doRequest(r)
- assert.Equal(t, http.StatusOK, w2.Code, "第 2 次请求应放行")
-
- w3 := doRequest(r)
- assert.Equal(t, http.StatusTooManyRequests, w3.Code, "第 3 次请求应被限流")
- }
-
- // TestUserModelRateLimit_ConcurrentMemory:内存路径并发安全验证。
- // RPM=2,并发 20 个请求,预期最多 2 个 200,其余 429。
- func TestUserModelRateLimit_ConcurrentMemory(t *testing.T) {
- setupTestDB(t)
- common.RedisEnabled = false
- setting.ModelRequestRateLimitEnabled = true
- setting.ModelRequestRateLimitDurationMinutes = 1
-
- require.NoError(t, model.DB.Create(&model.UserModelRateLimit{
- UserId: 30, Model: "deepseek-r1", Rpm: 2,
- }).Error)
-
- r := buildTestRouter(30, "deepseek-r1")
-
- var successCount int64
- var wg sync.WaitGroup
- totalRequests := 20
-
- for i := 0; i < totalRequests; i++ {
- wg.Add(1)
- go func() {
- defer wg.Done()
- w := doRequest(r)
- if w.Code == http.StatusOK {
- atomic.AddInt64(&successCount, 1)
- }
- }()
- }
- wg.Wait()
-
- sc := int(atomic.LoadInt64(&successCount))
- assert.LessOrEqual(t, sc, 2, "并发请求中最多 %d 个应成功,实际 %d 个", 2, sc)
- assert.GreaterOrEqual(t, sc, 1, "至少应有 1 个请求成功")
- }
-
- // TestUserModelRateLimit_UserIsolation:不同用户之间的限流互不影响
- func TestUserModelRateLimit_UserIsolation(t *testing.T) {
- setupTestDB(t)
- common.RedisEnabled = false
- setting.ModelRequestRateLimitEnabled = true
- setting.ModelRequestRateLimitDurationMinutes = 1
-
- // 两个用户都设置 RPM=2
- require.NoError(t, model.DB.Create(&model.UserModelRateLimit{
- UserId: 100, Model: "gpt-4", Rpm: 2,
- }).Error)
- require.NoError(t, model.DB.Create(&model.UserModelRateLimit{
- UserId: 101, Model: "gpt-4", Rpm: 2,
- }).Error)
-
- r1 := buildTestRouter(100, "gpt-4")
- r2 := buildTestRouter(101, "gpt-4")
-
- // 用户 100 用完配额
- w1 := doRequest(r1)
- assert.Equal(t, http.StatusOK, w1.Code)
- w2 := doRequest(r1)
- assert.Equal(t, http.StatusOK, w2.Code)
- w3 := doRequest(r1)
- assert.Equal(t, http.StatusTooManyRequests, w3.Code)
-
- // 用户 101 应该不受影响
- w4 := doRequest(r2)
- assert.Equal(t, http.StatusOK, w4.Code)
- w5 := doRequest(r2)
- assert.Equal(t, http.StatusOK, w5.Code)
- w6 := doRequest(r2)
- assert.Equal(t, http.StatusTooManyRequests, w6.Code)
- }
|