Вы не можете выбрать более 25 тем Темы должны начинаться с буквы или цифры, могут содержать дефисы(-) и должны содержать не более 35 символов.
 
 
 

265 строки
8.0 KiB

  1. package controller
  2. import (
  3. "bytes"
  4. "encoding/json"
  5. "net/http"
  6. "net/http/httptest"
  7. "strconv"
  8. "testing"
  9. "github.com/QuantumNous/new-api/common"
  10. "github.com/QuantumNous/new-api/model"
  11. "github.com/glebarez/sqlite"
  12. "github.com/gin-gonic/gin"
  13. "github.com/stretchr/testify/assert"
  14. "github.com/stretchr/testify/require"
  15. "gorm.io/gorm"
  16. )
  17. func setupEmailQuotaRuleControllerDB(t *testing.T) *gorm.DB {
  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. common.UsingSQLite = true
  26. common.RedisEnabled = false
  27. require.NoError(t, db.AutoMigrate(&model.EmailQuotaRule{}))
  28. t.Cleanup(func() {
  29. model.DB = origDB
  30. sqlDB.Close()
  31. })
  32. return db
  33. }
  34. func setupEmailQuotaRuleRouter() *gin.Engine {
  35. gin.SetMode(gin.TestMode)
  36. r := gin.New()
  37. g := r.Group("/api/email_quota_rule")
  38. {
  39. g.GET("/", GetAllEmailQuotaRules)
  40. g.POST("/", CreateEmailQuotaRule)
  41. g.PUT("/:id", UpdateEmailQuotaRule)
  42. g.DELETE("/:id", DeleteEmailQuotaRule)
  43. }
  44. return r
  45. }
  46. func TestGetAllEmailQuotaRules_Empty(t *testing.T) {
  47. setupEmailQuotaRuleControllerDB(t)
  48. router := setupEmailQuotaRuleRouter()
  49. w := httptest.NewRecorder()
  50. req, _ := http.NewRequest("GET", "/api/email_quota_rule/", nil)
  51. router.ServeHTTP(w, req)
  52. assert.Equal(t, http.StatusOK, w.Code)
  53. var resp map[string]interface{}
  54. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
  55. assert.True(t, resp["success"].(bool))
  56. assert.Empty(t, resp["data"])
  57. }
  58. func TestCreateEmailQuotaRule_Success(t *testing.T) {
  59. setupEmailQuotaRuleControllerDB(t)
  60. router := setupEmailQuotaRuleRouter()
  61. body, _ := json.Marshal(map[string]interface{}{
  62. "email_suffix": "@test.com",
  63. "quota": 500000,
  64. "description": "Test company",
  65. })
  66. w := httptest.NewRecorder()
  67. req, _ := http.NewRequest("POST", "/api/email_quota_rule/", bytes.NewReader(body))
  68. req.Header.Set("Content-Type", "application/json")
  69. router.ServeHTTP(w, req)
  70. assert.Equal(t, http.StatusOK, w.Code)
  71. var resp map[string]interface{}
  72. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
  73. assert.True(t, resp["success"].(bool))
  74. data := resp["data"].(map[string]interface{})
  75. assert.Equal(t, "@test.com", data["email_suffix"])
  76. assert.Equal(t, float64(500000), data["quota"])
  77. assert.Equal(t, true, data["enabled"])
  78. }
  79. func TestCreateEmailQuotaRule_Duplicate(t *testing.T) {
  80. setupEmailQuotaRuleControllerDB(t)
  81. router := setupEmailQuotaRuleRouter()
  82. body, _ := json.Marshal(map[string]interface{}{
  83. "email_suffix": "@dup.com",
  84. "quota": 100,
  85. })
  86. w := httptest.NewRecorder()
  87. req, _ := http.NewRequest("POST", "/api/email_quota_rule/", bytes.NewReader(body))
  88. req.Header.Set("Content-Type", "application/json")
  89. router.ServeHTTP(w, req)
  90. var firstResp map[string]interface{}
  91. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &firstResp))
  92. assert.True(t, firstResp["success"].(bool))
  93. // Second create should fail
  94. w2 := httptest.NewRecorder()
  95. req2, _ := http.NewRequest("POST", "/api/email_quota_rule/", bytes.NewReader(body))
  96. req2.Header.Set("Content-Type", "application/json")
  97. router.ServeHTTP(w2, req2)
  98. var resp2 map[string]interface{}
  99. require.NoError(t, json.Unmarshal(w2.Body.Bytes(), &resp2))
  100. assert.False(t, resp2["success"].(bool))
  101. }
  102. func TestCreateEmailQuotaRule_MissingFields(t *testing.T) {
  103. setupEmailQuotaRuleControllerDB(t)
  104. router := setupEmailQuotaRuleRouter()
  105. body, _ := json.Marshal(map[string]interface{}{
  106. "description": "no suffix",
  107. })
  108. w := httptest.NewRecorder()
  109. req, _ := http.NewRequest("POST", "/api/email_quota_rule/", bytes.NewReader(body))
  110. req.Header.Set("Content-Type", "application/json")
  111. router.ServeHTTP(w, req)
  112. var resp map[string]interface{}
  113. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
  114. assert.False(t, resp["success"].(bool))
  115. }
  116. func TestUpdateEmailQuotaRule_Success(t *testing.T) {
  117. setupEmailQuotaRuleControllerDB(t)
  118. router := setupEmailQuotaRuleRouter()
  119. // Create first
  120. rule := &model.EmailQuotaRule{EmailSuffix: "@up.com", Quota: 100, Enabled: true}
  121. require.NoError(t, rule.Insert())
  122. body, _ := json.Marshal(map[string]interface{}{
  123. "quota": 999,
  124. "description": "updated desc",
  125. })
  126. w := httptest.NewRecorder()
  127. req, _ := http.NewRequest("PUT", "/api/email_quota_rule/"+itoa(rule.Id), bytes.NewReader(body))
  128. req.Header.Set("Content-Type", "application/json")
  129. router.ServeHTTP(w, req)
  130. var resp map[string]interface{}
  131. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
  132. assert.True(t, resp["success"].(bool))
  133. data := resp["data"].(map[string]interface{})
  134. assert.Equal(t, float64(999), data["quota"])
  135. assert.Equal(t, "updated desc", data["description"])
  136. }
  137. func TestUpdateEmailQuotaRule_ToggleEnabled(t *testing.T) {
  138. setupEmailQuotaRuleControllerDB(t)
  139. router := setupEmailQuotaRuleRouter()
  140. rule := &model.EmailQuotaRule{EmailSuffix: "@toggle.com", Quota: 500, Enabled: true}
  141. require.NoError(t, rule.Insert())
  142. body, _ := json.Marshal(map[string]interface{}{
  143. "enabled": false,
  144. })
  145. w := httptest.NewRecorder()
  146. req, _ := http.NewRequest("PUT", "/api/email_quota_rule/"+itoa(rule.Id), bytes.NewReader(body))
  147. req.Header.Set("Content-Type", "application/json")
  148. router.ServeHTTP(w, req)
  149. var resp map[string]interface{}
  150. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
  151. assert.True(t, resp["success"].(bool))
  152. data := resp["data"].(map[string]interface{})
  153. assert.Equal(t, false, data["enabled"])
  154. // Cache should reflect disabled
  155. assert.Equal(t, int64(-1), model.MatchEmailQuotaRule("user@toggle.com"))
  156. }
  157. func TestDeleteEmailQuotaRule_Success(t *testing.T) {
  158. setupEmailQuotaRuleControllerDB(t)
  159. router := setupEmailQuotaRuleRouter()
  160. rule := &model.EmailQuotaRule{EmailSuffix: "@del.com", Quota: 100, Enabled: true}
  161. require.NoError(t, rule.Insert())
  162. w := httptest.NewRecorder()
  163. req, _ := http.NewRequest("DELETE", "/api/email_quota_rule/"+itoa(rule.Id), nil)
  164. router.ServeHTTP(w, req)
  165. var resp map[string]interface{}
  166. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
  167. assert.True(t, resp["success"].(bool))
  168. // Verify deleted
  169. assert.Equal(t, int64(-1), model.MatchEmailQuotaRule("user@del.com"))
  170. }
  171. func TestCRUD_FullFlow(t *testing.T) {
  172. setupEmailQuotaRuleControllerDB(t)
  173. router := setupEmailQuotaRuleRouter()
  174. // 1. Create
  175. body, _ := json.Marshal(map[string]interface{}{
  176. "email_suffix": "@full.com",
  177. "quota": 1000,
  178. "description": "full flow test",
  179. })
  180. w := httptest.NewRecorder()
  181. req, _ := http.NewRequest("POST", "/api/email_quota_rule/", bytes.NewReader(body))
  182. req.Header.Set("Content-Type", "application/json")
  183. router.ServeHTTP(w, req)
  184. var createResp map[string]interface{}
  185. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &createResp))
  186. assert.True(t, createResp["success"].(bool))
  187. // 2. List
  188. w = httptest.NewRecorder()
  189. req, _ = http.NewRequest("GET", "/api/email_quota_rule/", nil)
  190. router.ServeHTTP(w, req)
  191. var listResp map[string]interface{}
  192. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &listResp))
  193. data := listResp["data"].([]interface{})
  194. assert.Len(t, data, 1)
  195. // 3. Update
  196. ruleId := itoa(int(data[0].(map[string]interface{})["id"].(float64)))
  197. body, _ = json.Marshal(map[string]interface{}{"quota": 2000})
  198. w = httptest.NewRecorder()
  199. req, _ = http.NewRequest("PUT", "/api/email_quota_rule/"+ruleId, bytes.NewReader(body))
  200. req.Header.Set("Content-Type", "application/json")
  201. router.ServeHTTP(w, req)
  202. var updateResp map[string]interface{}
  203. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &updateResp))
  204. assert.True(t, updateResp["success"].(bool))
  205. // 4. Verify cache
  206. assert.Equal(t, int64(2000), model.MatchEmailQuotaRule("user@full.com"))
  207. // 5. Delete
  208. w = httptest.NewRecorder()
  209. req, _ = http.NewRequest("DELETE", "/api/email_quota_rule/"+ruleId, nil)
  210. router.ServeHTTP(w, req)
  211. var delResp map[string]interface{}
  212. require.NoError(t, json.Unmarshal(w.Body.Bytes(), &delResp))
  213. assert.True(t, delResp["success"].(bool))
  214. // 6. Verify cache cleared
  215. assert.Equal(t, int64(-1), model.MatchEmailQuotaRule("user@full.com"))
  216. }
  217. func itoa(i int) string {
  218. return strconv.Itoa(i)
  219. }