package controller import ( "bytes" "encoding/json" "net/http" "net/http/httptest" "strconv" "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 setupEmailQuotaRuleControllerDB(t *testing.T) *gorm.DB { 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.UsingSQLite = true common.RedisEnabled = false require.NoError(t, db.AutoMigrate(&model.EmailQuotaRule{})) t.Cleanup(func() { model.DB = origDB sqlDB.Close() }) return db } func setupEmailQuotaRuleRouter() *gin.Engine { gin.SetMode(gin.TestMode) r := gin.New() g := r.Group("/api/email_quota_rule") { g.GET("/", GetAllEmailQuotaRules) g.POST("/", CreateEmailQuotaRule) g.PUT("/:id", UpdateEmailQuotaRule) g.DELETE("/:id", DeleteEmailQuotaRule) } return r } func TestGetAllEmailQuotaRules_Empty(t *testing.T) { setupEmailQuotaRuleControllerDB(t) router := setupEmailQuotaRuleRouter() w := httptest.NewRecorder() req, _ := http.NewRequest("GET", "/api/email_quota_rule/", 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)) assert.Empty(t, resp["data"]) } func TestCreateEmailQuotaRule_Success(t *testing.T) { setupEmailQuotaRuleControllerDB(t) router := setupEmailQuotaRuleRouter() body, _ := json.Marshal(map[string]interface{}{ "email_suffix": "@test.com", "quota": 500000, "description": "Test company", }) w := httptest.NewRecorder() req, _ := http.NewRequest("POST", "/api/email_quota_rule/", 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)) data := resp["data"].(map[string]interface{}) assert.Equal(t, "@test.com", data["email_suffix"]) assert.Equal(t, float64(500000), data["quota"]) assert.Equal(t, true, data["enabled"]) } func TestCreateEmailQuotaRule_Duplicate(t *testing.T) { setupEmailQuotaRuleControllerDB(t) router := setupEmailQuotaRuleRouter() body, _ := json.Marshal(map[string]interface{}{ "email_suffix": "@dup.com", "quota": 100, }) w := httptest.NewRecorder() req, _ := http.NewRequest("POST", "/api/email_quota_rule/", bytes.NewReader(body)) req.Header.Set("Content-Type", "application/json") router.ServeHTTP(w, req) var firstResp map[string]interface{} require.NoError(t, json.Unmarshal(w.Body.Bytes(), &firstResp)) assert.True(t, firstResp["success"].(bool)) // Second create should fail w2 := httptest.NewRecorder() req2, _ := http.NewRequest("POST", "/api/email_quota_rule/", bytes.NewReader(body)) req2.Header.Set("Content-Type", "application/json") router.ServeHTTP(w2, req2) var resp2 map[string]interface{} require.NoError(t, json.Unmarshal(w2.Body.Bytes(), &resp2)) assert.False(t, resp2["success"].(bool)) } func TestCreateEmailQuotaRule_MissingFields(t *testing.T) { setupEmailQuotaRuleControllerDB(t) router := setupEmailQuotaRuleRouter() body, _ := json.Marshal(map[string]interface{}{ "description": "no suffix", }) w := httptest.NewRecorder() req, _ := http.NewRequest("POST", "/api/email_quota_rule/", bytes.NewReader(body)) req.Header.Set("Content-Type", "application/json") router.ServeHTTP(w, req) var resp map[string]interface{} require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) assert.False(t, resp["success"].(bool)) } func TestUpdateEmailQuotaRule_Success(t *testing.T) { setupEmailQuotaRuleControllerDB(t) router := setupEmailQuotaRuleRouter() // Create first rule := &model.EmailQuotaRule{EmailSuffix: "@up.com", Quota: 100, Enabled: true} require.NoError(t, rule.Insert()) body, _ := json.Marshal(map[string]interface{}{ "quota": 999, "description": "updated desc", }) w := httptest.NewRecorder() req, _ := http.NewRequest("PUT", "/api/email_quota_rule/"+itoa(rule.Id), bytes.NewReader(body)) req.Header.Set("Content-Type", "application/json") router.ServeHTTP(w, req) var resp map[string]interface{} require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) assert.True(t, resp["success"].(bool)) data := resp["data"].(map[string]interface{}) assert.Equal(t, float64(999), data["quota"]) assert.Equal(t, "updated desc", data["description"]) } func TestUpdateEmailQuotaRule_ToggleEnabled(t *testing.T) { setupEmailQuotaRuleControllerDB(t) router := setupEmailQuotaRuleRouter() rule := &model.EmailQuotaRule{EmailSuffix: "@toggle.com", Quota: 500, Enabled: true} require.NoError(t, rule.Insert()) body, _ := json.Marshal(map[string]interface{}{ "enabled": false, }) w := httptest.NewRecorder() req, _ := http.NewRequest("PUT", "/api/email_quota_rule/"+itoa(rule.Id), bytes.NewReader(body)) req.Header.Set("Content-Type", "application/json") router.ServeHTTP(w, req) var resp map[string]interface{} require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) assert.True(t, resp["success"].(bool)) data := resp["data"].(map[string]interface{}) assert.Equal(t, false, data["enabled"]) // Cache should reflect disabled assert.Equal(t, int64(-1), model.MatchEmailQuotaRule("user@toggle.com")) } func TestDeleteEmailQuotaRule_Success(t *testing.T) { setupEmailQuotaRuleControllerDB(t) router := setupEmailQuotaRuleRouter() rule := &model.EmailQuotaRule{EmailSuffix: "@del.com", Quota: 100, Enabled: true} require.NoError(t, rule.Insert()) w := httptest.NewRecorder() req, _ := http.NewRequest("DELETE", "/api/email_quota_rule/"+itoa(rule.Id), nil) router.ServeHTTP(w, req) var resp map[string]interface{} require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) assert.True(t, resp["success"].(bool)) // Verify deleted assert.Equal(t, int64(-1), model.MatchEmailQuotaRule("user@del.com")) } func TestCRUD_FullFlow(t *testing.T) { setupEmailQuotaRuleControllerDB(t) router := setupEmailQuotaRuleRouter() // 1. Create body, _ := json.Marshal(map[string]interface{}{ "email_suffix": "@full.com", "quota": 1000, "description": "full flow test", }) w := httptest.NewRecorder() req, _ := http.NewRequest("POST", "/api/email_quota_rule/", bytes.NewReader(body)) req.Header.Set("Content-Type", "application/json") router.ServeHTTP(w, req) var createResp map[string]interface{} require.NoError(t, json.Unmarshal(w.Body.Bytes(), &createResp)) assert.True(t, createResp["success"].(bool)) // 2. List w = httptest.NewRecorder() req, _ = http.NewRequest("GET", "/api/email_quota_rule/", nil) router.ServeHTTP(w, req) var listResp map[string]interface{} require.NoError(t, json.Unmarshal(w.Body.Bytes(), &listResp)) data := listResp["data"].([]interface{}) assert.Len(t, data, 1) // 3. Update ruleId := itoa(int(data[0].(map[string]interface{})["id"].(float64))) body, _ = json.Marshal(map[string]interface{}{"quota": 2000}) w = httptest.NewRecorder() req, _ = http.NewRequest("PUT", "/api/email_quota_rule/"+ruleId, bytes.NewReader(body)) req.Header.Set("Content-Type", "application/json") router.ServeHTTP(w, req) var updateResp map[string]interface{} require.NoError(t, json.Unmarshal(w.Body.Bytes(), &updateResp)) assert.True(t, updateResp["success"].(bool)) // 4. Verify cache assert.Equal(t, int64(2000), model.MatchEmailQuotaRule("user@full.com")) // 5. Delete w = httptest.NewRecorder() req, _ = http.NewRequest("DELETE", "/api/email_quota_rule/"+ruleId, nil) router.ServeHTTP(w, req) var delResp map[string]interface{} require.NoError(t, json.Unmarshal(w.Body.Bytes(), &delResp)) assert.True(t, delResp["success"].(bool)) // 6. Verify cache cleared assert.Equal(t, int64(-1), model.MatchEmailQuotaRule("user@full.com")) } func itoa(i int) string { return strconv.Itoa(i) }