package controller import ( "bytes" "encoding/json" "net/http" "net/http/httptest" "testing" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/model" "github.com/glebarez/sqlite" "github.com/gin-gonic/gin" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "gorm.io/gorm" ) func setupUserRateLimitControllerDB(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 common.RedisEnabled = false require.NoError(t, db.AutoMigrate(&model.UserModelRateLimit{})) t.Cleanup(func() { model.DB = origDB sqlDB.Close() }) } func setupUserRateLimitRouter() *gin.Engine { gin.SetMode(gin.TestMode) r := gin.New() g := r.Group("/api/user") { g.GET("/:id/rate-limits", GetUserRateLimits) g.PUT("/:id/rate-limits", SetUserRateLimits) } return r } func TestGetUserRateLimits_Empty(t *testing.T) { setupUserRateLimitControllerDB(t) router := setupUserRateLimitRouter() w := httptest.NewRecorder() req, _ := http.NewRequest(http.MethodGet, "/api/user/1/rate-limits", nil) router.ServeHTTP(w, req) assert.Equal(t, http.StatusOK, w.Code) var resp map[string]interface{} require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) assert.True(t, resp["success"].(bool)) // data 应为空列表(非 nil) data, ok := resp["data"] require.True(t, ok, "response should have data key") assert.Empty(t, data) } func TestSetUserRateLimits_OK(t *testing.T) { setupUserRateLimitControllerDB(t) router := setupUserRateLimitRouter() items := []model.UserModelRateLimit{ {Model: "gpt-4", Rpm: 60}, } body, _ := json.Marshal(items) w := httptest.NewRecorder() req, _ := http.NewRequest(http.MethodPut, "/api/user/1/rate-limits", bytes.NewReader(body)) req.Header.Set("Content-Type", "application/json") router.ServeHTTP(w, req) assert.Equal(t, http.StatusOK, w.Code) var resp map[string]interface{} require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) assert.True(t, resp["success"].(bool)) // 验证数据确实写入 list, err := model.GetUserModelRateLimits(1) require.NoError(t, err) require.Len(t, list, 1) assert.Equal(t, "gpt-4", list[0].Model) assert.Equal(t, 60, list[0].Rpm) } func TestSetUserRateLimits_NegativeRpm(t *testing.T) { setupUserRateLimitControllerDB(t) router := setupUserRateLimitRouter() items := []model.UserModelRateLimit{ {Model: "gpt-4", Rpm: -1}, } body, _ := json.Marshal(items) w := httptest.NewRecorder() req, _ := http.NewRequest(http.MethodPut, "/api/user/1/rate-limits", bytes.NewReader(body)) req.Header.Set("Content-Type", "application/json") router.ServeHTTP(w, req) assert.Equal(t, http.StatusOK, w.Code) var resp map[string]interface{} require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) assert.False(t, resp["success"].(bool)) } func TestGetUserRateLimits_InvalidId(t *testing.T) { setupUserRateLimitControllerDB(t) router := setupUserRateLimitRouter() w := httptest.NewRecorder() req, _ := http.NewRequest(http.MethodGet, "/api/user/notanid/rate-limits", nil) router.ServeHTTP(w, req) assert.Equal(t, http.StatusOK, w.Code) var resp map[string]interface{} require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) assert.False(t, resp["success"].(bool)) } func TestSetUserRateLimits_EmptySlice(t *testing.T) { setupUserRateLimitControllerDB(t) router := setupUserRateLimitRouter() // 先写一条 _ = model.SetUserModelRateLimits(2, []model.UserModelRateLimit{ {UserId: 2, Model: "claude-3-opus", Rpm: 10}, }) // 用空 slice 清空 body, _ := json.Marshal([]model.UserModelRateLimit{}) w := httptest.NewRecorder() req, _ := http.NewRequest(http.MethodPut, "/api/user/2/rate-limits", bytes.NewReader(body)) req.Header.Set("Content-Type", "application/json") router.ServeHTTP(w, req) assert.Equal(t, http.StatusOK, w.Code) var resp map[string]interface{} require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) assert.True(t, resp["success"].(bool)) list, err := model.GetUserModelRateLimits(2) require.NoError(t, err) assert.Empty(t, list) }