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")) }