package db import ( "context" "fmt" "sync" "testing" "github.com/stretchr/testify/require" "gorm.io/gorm" ) func createQuotaTestUser(t *testing.T, ctx context.Context, conn *gorm.DB, name string, dailyLimit int) *User { t.Helper() user := &User{Name: name, Domain: name, Email: name + "@example.test", Password: "test-only", DailyMessageLimit: dailyLimit} require.NoError(t, conn.WithContext(ctx).Create(user).Error) return user } func setFeatureQuota(t *testing.T, ctx context.Context, conn *gorm.DB, userID uint, feature string, limit int64) { t.Helper() require.NoError(t, NewFeatureControlsStore(conn).Save(ctx, userID, feature, UpdateFeatureControlOptions{Enabled: true, QuotaLimit: limit})) } func runConcurrently(n int, fn func(int) error) []error { start := make(chan struct{}) errs := make([]error, n) var wg sync.WaitGroup for i := range n { wg.Add(1) go func() { defer wg.Done() <-start errs[i] = fn(i) }() } close(start) wg.Wait() return errs } func requireOnlyQuotaErrors(t *testing.T, errs []error, quotaErr error) int { t.Helper() successes := 0 for _, err := range errs { if err == nil { successes++ continue } require.ErrorIs(t, err, quotaErr) } return successes } func TestQuestionQuotaConcurrentLastSlot(t *testing.T) { conn, _ := newTestDB(t) ctx := context.Background() owner := createQuotaTestUser(t, ctx, conn, "quota-question-owner", 0) setFeatureQuota(t, ctx, conn, owner.ID, FeatureQuestion, 1) questions := NewQuestionsStore(conn) errs := runConcurrently(12, func(i int) error { _, err := questions.Create(ctx, CreateQuestionOptions{UserID: owner.ID, Content: fmt.Sprintf("question-%d", i)}) return err }) require.Equal(t, 1, requireOnlyQuotaErrors(t, errs, ErrFeatureQuotaReached)) var count int64 require.NoError(t, conn.Model(&Question{}).Where("user_id = ?", owner.ID).Count(&count).Error) require.Equal(t, int64(1), count) fc, err := NewFeatureControlsStore(conn).Get(ctx, owner.ID, FeatureQuestion) require.NoError(t, err) require.Equal(t, int64(1), fc.QuotaCount) } func TestVoteQuotasConcurrentLastSlot(t *testing.T) { conn, _ := newTestDB(t) ctx := context.Background() owner := createQuotaTestUser(t, ctx, conn, "quota-vote-owner", 0) setFeatureQuota(t, ctx, conn, owner.ID, FeatureVote, 1) votes := NewVotesStore(conn) options := NewVoteOptionsStore(conn) submissions := NewVoteSubmissionsStore(conn) ids := make([]uint, 2) for i := range 2 { vote, err := votes.Create(ctx, CreateVoteOptions{UserID: owner.ID, Title: fmt.Sprintf("vote-%d", i), Type: VoteTypeSingle, QuotaLimit: 1}) require.NoError(t, err) require.NoError(t, votes.UpdateStatus(ctx, vote.ID, VoteStatusPublished)) option, err := options.Create(ctx, CreateVoteOptionOptions{VoteID: vote.ID, Content: "option"}) require.NoError(t, err) ids[i] = option.ID } errs := runConcurrently(12, func(i int) error { token := fmt.Sprintf("voter-%d", i) _, err := submissions.Create(ctx, CreateVoteSubmissionOptions{VoteID: uint(i%2 + 1), OptionID: ids[i%2], VoterToken: &token, QuotaUserID: owner.ID}) return err }) successes := 0 for _, err := range errs { if err == nil { successes++ } else { require.Error(t, err) } } require.Equal(t, 1, successes) var count int64 require.NoError(t, conn.Model(&VoteSubmission{}).Count(&count).Error) require.Equal(t, int64(1), count) fc, err := NewFeatureControlsStore(conn).Get(ctx, owner.ID, FeatureVote) require.NoError(t, err) require.Equal(t, int64(1), fc.QuotaCount) var ended int64 require.NoError(t, conn.Model(&Vote{}).Where("status = ?", VoteStatusEnded).Count(&ended).Error) require.Equal(t, int64(1), ended) } func TestDailyMessageLimitAcrossConversationPaths(t *testing.T) { conn, _ := newTestDB(t) ctx := context.Background() owner := createQuotaTestUser(t, ctx, conn, "quota-chat-owner", 1) asker := createQuotaTestUser(t, ctx, conn, "quota-chat-asker", 0) conversations := NewConversationStore(conn) existing, err := conversations.Create(ctx, CreateConversationOptions{OwnerID: owner.ID, AskerID: asker.ID, MaxMessages: 20}) require.NoError(t, err) var wg sync.WaitGroup start := make(chan struct{}) errs := make([]error, 2) wg.Add(2) go func() { defer wg.Done() <-start _, _, errs[0] = conversations.CreateWithFirstMessage(ctx, CreateConversationOptions{OwnerID: owner.ID, AskerID: asker.ID, MaxMessages: 20}, asker.ID, "first") }() go func() { defer wg.Done() <-start _, errs[1] = conversations.SendMessage(ctx, SendConversationMessageOptions{ConversationID: existing.ID, SenderID: asker.ID, Content: "existing"}) }() close(start) wg.Wait() require.Equal(t, 1, requireOnlyQuotaErrors(t, errs, ErrDailyMessageLimit)) var messages, convCount int64 require.NoError(t, conn.Model(&Message{}).Count(&messages).Error) require.NoError(t, conn.Model(&Conversation{}).Count(&convCount).Error) require.Equal(t, int64(1), messages) require.Equal(t, int64(1), convCount) } func TestOwnerReplyDoesNotConsumeAskerDailyQuota(t *testing.T) { conn, _ := newTestDB(t) ctx := context.Background() owner := createQuotaTestUser(t, ctx, conn, "quota-reply-owner", 1) asker := createQuotaTestUser(t, ctx, conn, "quota-reply-asker", 0) conversations := NewConversationStore(conn) conv, _, err := conversations.CreateWithFirstMessage(ctx, CreateConversationOptions{OwnerID: owner.ID, AskerID: asker.ID, MaxMessages: 20}, asker.ID, "first") require.NoError(t, err) _, err = conversations.SendMessage(ctx, SendConversationMessageOptions{ConversationID: conv.ID, SenderID: owner.ID, Content: "reply"}) require.NoError(t, err) _, err = conversations.SendMessage(ctx, SendConversationMessageOptions{ConversationID: conv.ID, SenderID: asker.ID, Content: "over"}) require.ErrorIs(t, err, ErrDailyMessageLimit) var count int64 require.NoError(t, conn.Model(&Message{}).Count(&count).Error) require.Equal(t, int64(2), count) }