|
- package model
-
- import (
- "testing"
-
- "github.com/QuantumNous/new-api/common"
- "github.com/glebarez/sqlite"
- "github.com/stretchr/testify/assert"
- "github.com/stretchr/testify/require"
- "gorm.io/gorm"
- )
-
- func setupEmailQuotaRuleDB(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 := DB
- DB = db
- common.UsingSQLite = true
- common.RedisEnabled = false
-
- require.NoError(t, db.AutoMigrate(&EmailQuotaRule{}))
-
- t.Cleanup(func() {
- DB = origDB
- sqlDB.Close()
- })
- return db
- }
-
- func TestMatchEmailQuotaRule_EmptyEmail(t *testing.T) {
- setupEmailQuotaRuleDB(t)
- LoadEmailQuotaCache()
-
- result := MatchEmailQuotaRule("")
- assert.Equal(t, int64(-1), result)
- }
-
- func TestMatchEmailQuotaRule_NoAtSign(t *testing.T) {
- setupEmailQuotaRuleDB(t)
- LoadEmailQuotaCache()
-
- result := MatchEmailQuotaRule("invalidemail")
- assert.Equal(t, int64(-1), result)
- }
-
- func TestMatchEmailQuotaRule_NoRules(t *testing.T) {
- setupEmailQuotaRuleDB(t)
- LoadEmailQuotaCache()
-
- result := MatchEmailQuotaRule("user@example.com")
- assert.Equal(t, int64(-1), result)
- }
-
- func TestMatchEmailQuotaRule_MatchEnabled(t *testing.T) {
- setupEmailQuotaRuleDB(t)
-
- rule := &EmailQuotaRule{
- EmailSuffix: "@example.com",
- Quota: 500000,
- Enabled: true,
- }
- require.NoError(t, rule.Insert())
-
- result := MatchEmailQuotaRule("user@example.com")
- assert.Equal(t, int64(500000), result)
- }
-
- func TestMatchEmailQuotaRule_CaseInsensitive(t *testing.T) {
- setupEmailQuotaRuleDB(t)
-
- rule := &EmailQuotaRule{
- EmailSuffix: "@Example.COM",
- Quota: 300000,
- Enabled: true,
- }
- require.NoError(t, rule.Insert())
-
- assert.Equal(t, int64(300000), MatchEmailQuotaRule("user@example.com"))
- assert.Equal(t, int64(300000), MatchEmailQuotaRule("user@EXAMPLE.COM"))
- assert.Equal(t, int64(300000), MatchEmailQuotaRule("user@Example.Com"))
- }
-
- func TestMatchEmailQuotaRule_DisabledRule(t *testing.T) {
- setupEmailQuotaRuleDB(t)
-
- rule := &EmailQuotaRule{
- EmailSuffix: "@disabled.com",
- Quota: 100000,
- Enabled: false,
- }
- require.NoError(t, rule.Insert())
-
- result := MatchEmailQuotaRule("user@disabled.com")
- assert.Equal(t, int64(-1), result)
- }
-
- func TestMatchEmailQuotaRule_NoMatch(t *testing.T) {
- setupEmailQuotaRuleDB(t)
-
- rule := &EmailQuotaRule{
- EmailSuffix: "@company.com",
- Quota: 500000,
- Enabled: true,
- }
- require.NoError(t, rule.Insert())
-
- result := MatchEmailQuotaRule("user@other.com")
- assert.Equal(t, int64(-1), result)
- }
-
- func TestEmailQuotaRule_Insert(t *testing.T) {
- setupEmailQuotaRuleDB(t)
-
- rule := &EmailQuotaRule{
- EmailSuffix: "@test.com",
- Quota: 100000,
- Enabled: true,
- Description: "Test rule",
- }
- require.NoError(t, rule.Insert())
- assert.Greater(t, rule.Id, 0)
- assert.Greater(t, rule.CreatedTime, int64(0))
- assert.Equal(t, rule.CreatedTime, rule.UpdatedTime)
-
- // Verify cache is populated
- assert.Equal(t, int64(100000), MatchEmailQuotaRule("user@test.com"))
- }
-
- func TestEmailQuotaRule_Insert_DuplicateSuffix(t *testing.T) {
- setupEmailQuotaRuleDB(t)
-
- rule1 := &EmailQuotaRule{EmailSuffix: "@dup.com", Quota: 100, Enabled: true}
- require.NoError(t, rule1.Insert())
-
- rule2 := &EmailQuotaRule{EmailSuffix: "@dup.com", Quota: 200, Enabled: true}
- assert.Error(t, rule2.Insert())
- }
-
- func TestEmailQuotaRule_Update(t *testing.T) {
- setupEmailQuotaRuleDB(t)
-
- rule := &EmailQuotaRule{EmailSuffix: "@update.com", Quota: 100, Enabled: true}
- require.NoError(t, rule.Insert())
-
- rule.Quota = 999
- rule.Description = "updated"
- require.NoError(t, rule.Update())
-
- // Verify DB
- found, err := GetEmailQuotaRuleById(rule.Id)
- require.NoError(t, err)
- assert.Equal(t, int64(999), found.Quota)
- assert.Equal(t, "updated", found.Description)
-
- // Verify cache refreshed
- assert.Equal(t, int64(999), MatchEmailQuotaRule("user@update.com"))
- }
-
- func TestEmailQuotaRule_Update_Disable(t *testing.T) {
- setupEmailQuotaRuleDB(t)
-
- rule := &EmailQuotaRule{EmailSuffix: "@toggled.com", Quota: 500, Enabled: true}
- require.NoError(t, rule.Insert())
- assert.Equal(t, int64(500), MatchEmailQuotaRule("user@toggled.com"))
-
- rule.Enabled = false
- require.NoError(t, rule.Update())
-
- assert.Equal(t, int64(-1), MatchEmailQuotaRule("user@toggled.com"))
- }
-
- func TestEmailQuotaRule_Delete(t *testing.T) {
- setupEmailQuotaRuleDB(t)
-
- rule := &EmailQuotaRule{EmailSuffix: "@delete.com", Quota: 100, Enabled: true}
- require.NoError(t, rule.Insert())
- assert.Equal(t, int64(100), MatchEmailQuotaRule("user@delete.com"))
-
- require.NoError(t, rule.Delete())
- assert.Equal(t, int64(-1), MatchEmailQuotaRule("user@delete.com"))
- }
-
- func TestGetAllEmailQuotaRules(t *testing.T) {
- setupEmailQuotaRuleDB(t)
-
- r1 := &EmailQuotaRule{EmailSuffix: "@a.com", Quota: 100, Enabled: true}
- r2 := &EmailQuotaRule{EmailSuffix: "@b.com", Quota: 200, Enabled: false}
- require.NoError(t, r1.Insert())
- require.NoError(t, r2.Insert())
-
- list, err := GetAllEmailQuotaRules()
- require.NoError(t, err)
- assert.Len(t, list, 2)
- // Ordered by id ASC
- assert.Equal(t, "@a.com", list[0].EmailSuffix)
- assert.Equal(t, "@b.com", list[1].EmailSuffix)
- }
-
- func TestGetEmailQuotaRuleBySuffix(t *testing.T) {
- setupEmailQuotaRuleDB(t)
-
- rule := &EmailQuotaRule{EmailSuffix: "@find.com", Quota: 300, Enabled: true}
- require.NoError(t, rule.Insert())
-
- found, err := GetEmailQuotaRuleBySuffix("@find.com")
- require.NoError(t, err)
- assert.Equal(t, int64(300), found.Quota)
-
- _, err = GetEmailQuotaRuleBySuffix("@notexist.com")
- assert.Error(t, err)
- }
-
- func TestLoadEmailQuotaCache_OnlyEnabled(t *testing.T) {
- setupEmailQuotaRuleDB(t)
-
- r1 := &EmailQuotaRule{EmailSuffix: "@enabled.com", Quota: 100, Enabled: true}
- r2 := &EmailQuotaRule{EmailSuffix: "@disabled.com", Quota: 200, Enabled: false}
- require.NoError(t, r1.Insert())
- require.NoError(t, r2.Insert())
- LoadEmailQuotaCache()
-
- assert.Equal(t, int64(100), MatchEmailQuotaRule("user@enabled.com"))
- assert.Equal(t, int64(-1), MatchEmailQuotaRule("user@disabled.com"))
- }
-
- func TestMatchEmailQuotaRule_MultipleRules(t *testing.T) {
- setupEmailQuotaRuleDB(t)
-
- rules := []*EmailQuotaRule{
- {EmailSuffix: "@company.com", Quota: 500000, Enabled: true},
- {EmailSuffix: "@tsinghua.edu.cn", Quota: 1000000, Enabled: true},
- {EmailSuffix: "@vip.org", Quota: 2000000, Enabled: true},
- }
- for _, r := range rules {
- require.NoError(t, r.Insert())
- }
-
- assert.Equal(t, int64(500000), MatchEmailQuotaRule("user@company.com"))
- assert.Equal(t, int64(1000000), MatchEmailQuotaRule("student@tsinghua.edu.cn"))
- assert.Equal(t, int64(2000000), MatchEmailQuotaRule("admin@vip.org"))
- assert.Equal(t, int64(-1), MatchEmailQuotaRule("random@unknown.net"))
- }
|