tamabox/internal/db/quota_concurrency_test.go

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