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