Não pode escolher mais do que 25 tópicos Os tópicos devem começar com uma letra ou um número, podem incluir traços ('-') e podem ter até 35 caracteres.
 
 
 

247 linhas
6.5 KiB

  1. package model
  2. import (
  3. "testing"
  4. "github.com/QuantumNous/new-api/common"
  5. "github.com/glebarez/sqlite"
  6. "github.com/stretchr/testify/assert"
  7. "github.com/stretchr/testify/require"
  8. "gorm.io/gorm"
  9. )
  10. func setupEmailQuotaRuleDB(t *testing.T) *gorm.DB {
  11. t.Helper()
  12. db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
  13. require.NoError(t, err)
  14. sqlDB, _ := db.DB()
  15. sqlDB.SetMaxOpenConns(1)
  16. origDB := DB
  17. DB = db
  18. common.UsingSQLite = true
  19. common.RedisEnabled = false
  20. require.NoError(t, db.AutoMigrate(&EmailQuotaRule{}))
  21. t.Cleanup(func() {
  22. DB = origDB
  23. sqlDB.Close()
  24. })
  25. return db
  26. }
  27. func TestMatchEmailQuotaRule_EmptyEmail(t *testing.T) {
  28. setupEmailQuotaRuleDB(t)
  29. LoadEmailQuotaCache()
  30. result := MatchEmailQuotaRule("")
  31. assert.Equal(t, int64(-1), result)
  32. }
  33. func TestMatchEmailQuotaRule_NoAtSign(t *testing.T) {
  34. setupEmailQuotaRuleDB(t)
  35. LoadEmailQuotaCache()
  36. result := MatchEmailQuotaRule("invalidemail")
  37. assert.Equal(t, int64(-1), result)
  38. }
  39. func TestMatchEmailQuotaRule_NoRules(t *testing.T) {
  40. setupEmailQuotaRuleDB(t)
  41. LoadEmailQuotaCache()
  42. result := MatchEmailQuotaRule("user@example.com")
  43. assert.Equal(t, int64(-1), result)
  44. }
  45. func TestMatchEmailQuotaRule_MatchEnabled(t *testing.T) {
  46. setupEmailQuotaRuleDB(t)
  47. rule := &EmailQuotaRule{
  48. EmailSuffix: "@example.com",
  49. Quota: 500000,
  50. Enabled: true,
  51. }
  52. require.NoError(t, rule.Insert())
  53. result := MatchEmailQuotaRule("user@example.com")
  54. assert.Equal(t, int64(500000), result)
  55. }
  56. func TestMatchEmailQuotaRule_CaseInsensitive(t *testing.T) {
  57. setupEmailQuotaRuleDB(t)
  58. rule := &EmailQuotaRule{
  59. EmailSuffix: "@Example.COM",
  60. Quota: 300000,
  61. Enabled: true,
  62. }
  63. require.NoError(t, rule.Insert())
  64. assert.Equal(t, int64(300000), MatchEmailQuotaRule("user@example.com"))
  65. assert.Equal(t, int64(300000), MatchEmailQuotaRule("user@EXAMPLE.COM"))
  66. assert.Equal(t, int64(300000), MatchEmailQuotaRule("user@Example.Com"))
  67. }
  68. func TestMatchEmailQuotaRule_DisabledRule(t *testing.T) {
  69. setupEmailQuotaRuleDB(t)
  70. rule := &EmailQuotaRule{
  71. EmailSuffix: "@disabled.com",
  72. Quota: 100000,
  73. Enabled: false,
  74. }
  75. require.NoError(t, rule.Insert())
  76. result := MatchEmailQuotaRule("user@disabled.com")
  77. assert.Equal(t, int64(-1), result)
  78. }
  79. func TestMatchEmailQuotaRule_NoMatch(t *testing.T) {
  80. setupEmailQuotaRuleDB(t)
  81. rule := &EmailQuotaRule{
  82. EmailSuffix: "@company.com",
  83. Quota: 500000,
  84. Enabled: true,
  85. }
  86. require.NoError(t, rule.Insert())
  87. result := MatchEmailQuotaRule("user@other.com")
  88. assert.Equal(t, int64(-1), result)
  89. }
  90. func TestEmailQuotaRule_Insert(t *testing.T) {
  91. setupEmailQuotaRuleDB(t)
  92. rule := &EmailQuotaRule{
  93. EmailSuffix: "@test.com",
  94. Quota: 100000,
  95. Enabled: true,
  96. Description: "Test rule",
  97. }
  98. require.NoError(t, rule.Insert())
  99. assert.Greater(t, rule.Id, 0)
  100. assert.Greater(t, rule.CreatedTime, int64(0))
  101. assert.Equal(t, rule.CreatedTime, rule.UpdatedTime)
  102. // Verify cache is populated
  103. assert.Equal(t, int64(100000), MatchEmailQuotaRule("user@test.com"))
  104. }
  105. func TestEmailQuotaRule_Insert_DuplicateSuffix(t *testing.T) {
  106. setupEmailQuotaRuleDB(t)
  107. rule1 := &EmailQuotaRule{EmailSuffix: "@dup.com", Quota: 100, Enabled: true}
  108. require.NoError(t, rule1.Insert())
  109. rule2 := &EmailQuotaRule{EmailSuffix: "@dup.com", Quota: 200, Enabled: true}
  110. assert.Error(t, rule2.Insert())
  111. }
  112. func TestEmailQuotaRule_Update(t *testing.T) {
  113. setupEmailQuotaRuleDB(t)
  114. rule := &EmailQuotaRule{EmailSuffix: "@update.com", Quota: 100, Enabled: true}
  115. require.NoError(t, rule.Insert())
  116. rule.Quota = 999
  117. rule.Description = "updated"
  118. require.NoError(t, rule.Update())
  119. // Verify DB
  120. found, err := GetEmailQuotaRuleById(rule.Id)
  121. require.NoError(t, err)
  122. assert.Equal(t, int64(999), found.Quota)
  123. assert.Equal(t, "updated", found.Description)
  124. // Verify cache refreshed
  125. assert.Equal(t, int64(999), MatchEmailQuotaRule("user@update.com"))
  126. }
  127. func TestEmailQuotaRule_Update_Disable(t *testing.T) {
  128. setupEmailQuotaRuleDB(t)
  129. rule := &EmailQuotaRule{EmailSuffix: "@toggled.com", Quota: 500, Enabled: true}
  130. require.NoError(t, rule.Insert())
  131. assert.Equal(t, int64(500), MatchEmailQuotaRule("user@toggled.com"))
  132. rule.Enabled = false
  133. require.NoError(t, rule.Update())
  134. assert.Equal(t, int64(-1), MatchEmailQuotaRule("user@toggled.com"))
  135. }
  136. func TestEmailQuotaRule_Delete(t *testing.T) {
  137. setupEmailQuotaRuleDB(t)
  138. rule := &EmailQuotaRule{EmailSuffix: "@delete.com", Quota: 100, Enabled: true}
  139. require.NoError(t, rule.Insert())
  140. assert.Equal(t, int64(100), MatchEmailQuotaRule("user@delete.com"))
  141. require.NoError(t, rule.Delete())
  142. assert.Equal(t, int64(-1), MatchEmailQuotaRule("user@delete.com"))
  143. }
  144. func TestGetAllEmailQuotaRules(t *testing.T) {
  145. setupEmailQuotaRuleDB(t)
  146. r1 := &EmailQuotaRule{EmailSuffix: "@a.com", Quota: 100, Enabled: true}
  147. r2 := &EmailQuotaRule{EmailSuffix: "@b.com", Quota: 200, Enabled: false}
  148. require.NoError(t, r1.Insert())
  149. require.NoError(t, r2.Insert())
  150. list, err := GetAllEmailQuotaRules()
  151. require.NoError(t, err)
  152. assert.Len(t, list, 2)
  153. // Ordered by id ASC
  154. assert.Equal(t, "@a.com", list[0].EmailSuffix)
  155. assert.Equal(t, "@b.com", list[1].EmailSuffix)
  156. }
  157. func TestGetEmailQuotaRuleBySuffix(t *testing.T) {
  158. setupEmailQuotaRuleDB(t)
  159. rule := &EmailQuotaRule{EmailSuffix: "@find.com", Quota: 300, Enabled: true}
  160. require.NoError(t, rule.Insert())
  161. found, err := GetEmailQuotaRuleBySuffix("@find.com")
  162. require.NoError(t, err)
  163. assert.Equal(t, int64(300), found.Quota)
  164. _, err = GetEmailQuotaRuleBySuffix("@notexist.com")
  165. assert.Error(t, err)
  166. }
  167. func TestLoadEmailQuotaCache_OnlyEnabled(t *testing.T) {
  168. setupEmailQuotaRuleDB(t)
  169. r1 := &EmailQuotaRule{EmailSuffix: "@enabled.com", Quota: 100, Enabled: true}
  170. r2 := &EmailQuotaRule{EmailSuffix: "@disabled.com", Quota: 200, Enabled: false}
  171. require.NoError(t, r1.Insert())
  172. require.NoError(t, r2.Insert())
  173. LoadEmailQuotaCache()
  174. assert.Equal(t, int64(100), MatchEmailQuotaRule("user@enabled.com"))
  175. assert.Equal(t, int64(-1), MatchEmailQuotaRule("user@disabled.com"))
  176. }
  177. func TestMatchEmailQuotaRule_MultipleRules(t *testing.T) {
  178. setupEmailQuotaRuleDB(t)
  179. rules := []*EmailQuotaRule{
  180. {EmailSuffix: "@company.com", Quota: 500000, Enabled: true},
  181. {EmailSuffix: "@tsinghua.edu.cn", Quota: 1000000, Enabled: true},
  182. {EmailSuffix: "@vip.org", Quota: 2000000, Enabled: true},
  183. }
  184. for _, r := range rules {
  185. require.NoError(t, r.Insert())
  186. }
  187. assert.Equal(t, int64(500000), MatchEmailQuotaRule("user@company.com"))
  188. assert.Equal(t, int64(1000000), MatchEmailQuotaRule("student@tsinghua.edu.cn"))
  189. assert.Equal(t, int64(2000000), MatchEmailQuotaRule("admin@vip.org"))
  190. assert.Equal(t, int64(-1), MatchEmailQuotaRule("random@unknown.net"))
  191. }