thread-master/apps/backend/internal/module/usage/usecase/service_test.go

package usecase_test

import (
	"context"
	"fmt"
	"sync"
	"testing"
	"time"

	"apps/backend/internal/module/usage/domain"
	"apps/backend/internal/module/usage/repository"
	"apps/backend/internal/module/usage/usecase"

	"github.com/stretchr/testify/require"
)

func newUsage(uid int64, mode string) *usecase.Service {
	r := repository.NewMemory()
	key := func(m string) string { return fmt.Sprintf("%d:%s", uid, m) }
	res := &usecase.StaticResolver{Map: map[string]string{
		key(domain.MeterAICopy):     mode,
		key(domain.MeterAIResearch): mode,
		key(domain.MeterWebSearch):  mode,
		key(domain.MeterAIImage):    mode,
	}}
	// also allow mixed tests
	if mode == "mixed" {
		res.Map = map[string]string{
			key(domain.MeterAICopy):     domain.KeyModeByok,
			key(domain.MeterAIResearch): domain.KeyModeByok,
			key(domain.MeterWebSearch):  domain.KeyModePlatform,
			key(domain.MeterAIImage):    domain.KeyModeByok,
		}
	}
	return usecase.New(r, res)
}

func TestUSG_01_NewMemberSummarySplit(t *testing.T) {
	svc := newUsage(1_000_001, domain.KeyModePlatform)
	sum, err := svc.GetSummary(context.Background(), 1_000_001, "")
	require.NoError(t, err)
	require.Equal(t, domain.PlanFree, sum.PlanID)
	require.Equal(t, 0, sum.Platform.CreditsUsed)
	require.Equal(t, 0, sum.Byok.CallCount)
	require.Greater(t, sum.Platform.CreditsRemaining, 0)
}

func TestUSG_02_PlatformOnly(t *testing.T) {
	svc := newUsage(1_000_001, domain.KeyModePlatform)
	_, err := svc.RecordCall(context.Background(), 1_000_001, domain.MeterAICopy, domain.KeyModePlatform, "copy", "test")
	require.NoError(t, err)
	sum, err := svc.GetSummary(context.Background(), 1_000_001, domain.CurrentMonthKey())
	require.NoError(t, err)
	require.Equal(t, 1, sum.Platform.CreditsUsed)
	require.Equal(t, 0, sum.Byok.CallCount)
}

func TestUSG_03_ByokOnly(t *testing.T) {
	svc := newUsage(1_000_001, domain.KeyModeByok)
	before, _ := svc.GetSummary(context.Background(), 1_000_001, "")
	_, err := svc.RecordCall(context.Background(), 1_000_001, domain.MeterAICopy, domain.KeyModeByok, "copy", "test")
	require.NoError(t, err)
	sum, err := svc.GetSummary(context.Background(), 1_000_001, domain.CurrentMonthKey())
	require.NoError(t, err)
	require.Equal(t, before.Platform.CreditsRemaining, sum.Platform.CreditsRemaining)
	require.Equal(t, 1, sum.Byok.CallCount)
	require.Equal(t, 0, sum.Platform.CreditsUsed)
}

func TestUSG_04_MixedAIByokSearchPlatform(t *testing.T) {
	svc := newUsage(1_000_001, "mixed")
	_, _ = svc.RecordCall(context.Background(), 1_000_001, domain.MeterAICopy, domain.KeyModeByok, "ai", "t")
	_, _ = svc.RecordCall(context.Background(), 1_000_001, domain.MeterWebSearch, domain.KeyModePlatform, "search", "t")
	sum, err := svc.GetSummary(context.Background(), 1_000_001, domain.CurrentMonthKey())
	require.NoError(t, err)
	require.Equal(t, 1, sum.Platform.CreditsUsed) // search only
	require.Equal(t, 1, sum.Byok.CallCount)       // ai only
}

func TestUSG_05_ListEventsHasKeyMode(t *testing.T) {
	svc := newUsage(1_000_001, domain.KeyModePlatform)
	_, _ = svc.RecordCall(context.Background(), 1_000_001, domain.MeterAICopy, domain.KeyModePlatform, "x", "t")
	list, err := svc.ListEvents(context.Background(), 1_000_001, "", "all", 10)
	require.NoError(t, err)
	require.NotEmpty(t, list)
	require.Equal(t, domain.KeyModePlatform, list[0].KeyMode)
	require.NotEmpty(t, list[0].Meter)
}

func TestUSG_06_FilterByok(t *testing.T) {
	svc := newUsage(1_000_001, "mixed")
	_, _ = svc.RecordCall(context.Background(), 1_000_001, domain.MeterAICopy, domain.KeyModeByok, "a", "t")
	_, _ = svc.RecordCall(context.Background(), 1_000_001, domain.MeterWebSearch, domain.KeyModePlatform, "s", "t")
	list, err := svc.ListEvents(context.Background(), 1_000_001, "", domain.KeyModeByok, 50)
	require.NoError(t, err)
	require.Len(t, list, 1)
	require.Equal(t, domain.KeyModeByok, list[0].KeyMode)
}

func TestUSG_07_FilterPlatform(t *testing.T) {
	svc := newUsage(1_000_001, "mixed")
	_, _ = svc.RecordCall(context.Background(), 1_000_001, domain.MeterAICopy, domain.KeyModeByok, "a", "t")
	_, _ = svc.RecordCall(context.Background(), 1_000_001, domain.MeterWebSearch, domain.KeyModePlatform, "s", "t")
	list, err := svc.ListEvents(context.Background(), 1_000_001, "", domain.KeyModePlatform, 50)
	require.NoError(t, err)
	require.Len(t, list, 1)
	require.Equal(t, domain.KeyModePlatform, list[0].KeyMode)
}

func TestUSG_08_NoMergedTotalOnly(t *testing.T) {
	svc := newUsage(1_000_001, "mixed")
	_, _ = svc.RecordCall(context.Background(), 1_000_001, domain.MeterAICopy, domain.KeyModeByok, "a", "t")
	_, _ = svc.RecordCall(context.Background(), 1_000_001, domain.MeterWebSearch, domain.KeyModePlatform, "s", "t")
	sum, err := svc.GetSummary(context.Background(), 1_000_001, domain.CurrentMonthKey())
	require.NoError(t, err)
	// Must have separate blocks; platform remaining is not reduced by byok
	require.NotNil(t, sum.Platform)
	require.NotNil(t, sum.Byok)
	require.Equal(t, 1, sum.Platform.CallCount)
	require.Equal(t, 1, sum.Byok.CallCount)
	// total_credits is plan allotment (platform), not call merge
	require.Equal(t, sum.Platform.CreditsTotal, sum.TotalCredits)
}

func TestST_07_ByokNoCreditCharge(t *testing.T) {
	svc := newUsage(1_000_001, domain.KeyModeByok)
	before, _ := svc.GetSummary(context.Background(), 1_000_001, "")
	mode, err := svc.PrepareCall(context.Background(), 1_000_001, domain.MeterAICopy)
	require.NoError(t, err)
	require.Equal(t, domain.KeyModeByok, mode)
	_, err = svc.RecordCall(context.Background(), 1_000_001, domain.MeterAICopy, mode, "byok", "ai")
	require.NoError(t, err)
	after, _ := svc.GetSummary(context.Background(), 1_000_001, domain.CurrentMonthKey())
	require.Equal(t, before.Platform.CreditsRemaining, after.Platform.CreditsRemaining)
	require.Equal(t, 1, after.Byok.CallCount)
}

func TestST_09_NoKeyFails(t *testing.T) {
	r := repository.NewMemory()
	svc := usecase.New(r, &usecase.StaticResolver{Map: map[string]string{}})
	_, err := svc.PrepareCall(context.Background(), 1_000_001, domain.MeterAICopy)
	require.ErrorIs(t, err, domain.ErrNoKey)
}

func TestST_13_PlatformQuota(t *testing.T) {
	svc := newUsage(1_000_001, domain.KeyModePlatform)
	// free: ai_copy soft cap 60pt；先燒滿分項
	for i := 0; i < 60; i++ {
		mode, err := svc.PrepareCall(context.Background(), 1_000_001, domain.MeterAICopy)
		require.NoError(t, err)
		_, err = svc.RecordCall(context.Background(), 1_000_001, domain.MeterAICopy, mode, "x", "t")
		require.NoError(t, err)
	}
	_, err := svc.PrepareCall(context.Background(), 1_000_001, domain.MeterAICopy)
	require.ErrorIs(t, err, domain.ErrMeterCapExceeded)
	// 分項鎖死後，其他 meter 在總池仍有餘額時可繼續
	_, err = svc.PrepareCall(context.Background(), 1_000_001, domain.MeterWebSearch)
	require.NoError(t, err)
	_, err = svc.RecordCall(context.Background(), 1_000_001, domain.MeterWebSearch, domain.KeyModePlatform, "s", "t")
	require.NoError(t, err)
	// 燒光剩餘總池：search 30 + research 15 + image 15 = 60
	for i := 1; i < 30; i++ {
		mode, err := svc.PrepareCall(context.Background(), 1_000_001, domain.MeterWebSearch)
		require.NoError(t, err)
		_, err = svc.RecordCall(context.Background(), 1_000_001, domain.MeterWebSearch, mode, "s", "t")
		require.NoError(t, err)
	}
	for i := 0; i < 5; i++ { // 5*3=15
		mode, err := svc.PrepareCall(context.Background(), 1_000_001, domain.MeterAIResearch)
		require.NoError(t, err)
		_, err = svc.RecordCall(context.Background(), 1_000_001, domain.MeterAIResearch, mode, "r", "t")
		require.NoError(t, err)
	}
	// image 15pt = 3 張
	for i := 0; i < 3; i++ {
		mode, err := svc.PrepareCall(context.Background(), 1_000_001, domain.MeterAIImage)
		require.NoError(t, err)
		_, err = svc.RecordCall(context.Background(), 1_000_001, domain.MeterAIImage, mode, "img", "t")
		require.NoError(t, err)
	}
	_, err = svc.PrepareCall(context.Background(), 1_000_001, domain.MeterWebSearch)
	require.ErrorIs(t, err, domain.ErrQuotaExceeded)
}

func TestPlatformReservation_ConcurrentOneCreditCannotBeOverspent(t *testing.T) {
	repo := repository.NewMemory()
	const attempts = 32
	start := make(chan struct{})
	results := make(chan error, attempts)
	var wg sync.WaitGroup
	for i := 0; i < attempts; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			<-start
			results <- repo.ReservePlatform(context.Background(), 1_000_003, domain.CurrentMonthKey(), domain.MeterAICopy, 1, 1, 1)
		}()
	}
	close(start)
	wg.Wait()
	close(results)

	succeeded := 0
	for err := range results {
		if err == nil {
			succeeded++
			continue
		}
		require.ErrorIs(t, err, domain.ErrQuotaExceeded)
	}
	require.Equal(t, 1, succeeded)
	counter, err := repo.GetMonthlyCounter(context.Background(), 1_000_003, domain.CurrentMonthKey())
	require.NoError(t, err)
	require.Equal(t, 1, counter.TotalCredits)
	require.Equal(t, 1, counter.ByMeter[domain.MeterAICopy])
}

func TestPrepareCall_BYOKDoesNotCreatePlatformCounter(t *testing.T) {
	repo := repository.NewMemory()
	uid := int64(1_000_004)
	svc := usecase.New(repo, &usecase.StaticResolver{Map: map[string]string{
		fmt.Sprintf("%d:%s", uid, domain.MeterAICopy): domain.KeyModeByok,
	}})
	mode, err := svc.PrepareCall(context.Background(), uid, domain.MeterAICopy)
	require.NoError(t, err)
	require.Equal(t, domain.KeyModeByok, mode)
	_, err = repo.GetMonthlyCounter(context.Background(), uid, domain.CurrentMonthKey())
	require.ErrorIs(t, err, domain.ErrNotFound)
}

func TestUSG_MeterCap_Image(t *testing.T) {
	svc := newUsage(1_000_001, domain.KeyModePlatform)
	// free ai_image cap 15pt = 3 次；第 4 次因分項上限擋
	for i := 0; i < 3; i++ {
		_, err := svc.RecordCall(context.Background(), 1_000_001, domain.MeterAIImage, domain.KeyModePlatform, "img", "t")
		require.NoError(t, err)
	}
	_, err := svc.PrepareCall(context.Background(), 1_000_001, domain.MeterAIImage)
	require.ErrorIs(t, err, domain.ErrMeterCapExceeded)
}

func TestUSG_PlatformCapacity_BlocksPlatformOnly(t *testing.T) {
	// platform user blocked when gate cools down
	plat := newUsage(1_000_001, domain.KeyModePlatform)
	plat.Gate = usecase.NewMemoryPlatformGate(0, 0)
	plat.MarkPlatformLimited(time.Now().Add(time.Hour))
	_, err := plat.PrepareCall(context.Background(), 1_000_001, domain.MeterAICopy)
	require.ErrorIs(t, err, domain.ErrPlatformCapacity)

	// BYOK still works (gate not consulted)
	byok := newUsage(1_000_002, domain.KeyModeByok)
	byok.Gate = usecase.NewMemoryPlatformGate(0, 0)
	byok.MarkPlatformLimited(time.Now().Add(time.Hour))
	mode, err := byok.PrepareCall(context.Background(), 1_000_002, domain.MeterAICopy)
	require.NoError(t, err)
	require.Equal(t, domain.KeyModeByok, mode)
}

func TestUSG_PlatformCapacity_RPM(t *testing.T) {
	svc := newUsage(1_000_001, domain.KeyModePlatform)
	svc.Gate = usecase.NewMemoryPlatformGate(2, 0) // 2 AI / min
	_, err := svc.PrepareCall(context.Background(), 1_000_001, domain.MeterAICopy)
	require.NoError(t, err)
	_, err = svc.PrepareCall(context.Background(), 1_000_001, domain.MeterAICopy)
	require.NoError(t, err)
	_, err = svc.PrepareCall(context.Background(), 1_000_001, domain.MeterAICopy)
	require.ErrorIs(t, err, domain.ErrPlatformCapacity)
}

func TestUSG_11_LegacyPurchaseCannotActivatePaidPlan(t *testing.T) {
	svc := newUsage(1_000_001, domain.KeyModePlatform)
	p, err := svc.PurchasePlan(context.Background(), 1_000_001, domain.PlanStarter, "mock-pay-1")
	require.Nil(t, p)
	require.ErrorIs(t, err, domain.ErrPurchaseDisabled)
	prefs, _ := svc.GetPrefs(context.Background(), 1_000_001)
	require.Equal(t, domain.PlanFree, prefs.PlanID)
	list, err := svc.ListMyPurchases(context.Background(), 1_000_001, 10)
	require.NoError(t, err)
	require.Empty(t, list)
}

func TestUSG_10_TenantAnalyticsSplit(t *testing.T) {
	svc := newUsage(1_000_001, "mixed")
	_, _ = svc.RecordCall(context.Background(), 1_000_001, domain.MeterAICopy, domain.KeyModeByok, "a", "t")
	_, _ = svc.RecordCall(context.Background(), 1_000_001, domain.MeterWebSearch, domain.KeyModePlatform, "s", "t")
	ten, err := svc.TenantSummary(context.Background(), domain.CurrentMonthKey())
	require.NoError(t, err)
	require.Equal(t, 1, ten.PlatformCreditsUsed)
	require.Equal(t, 1, ten.ByokCallCount)
}

func TestUSG_14_MemberCannotSetOthersPrefs(t *testing.T) {
	svc := newUsage(1_000_001, domain.KeyModePlatform)
	_, err := svc.SetPrefs(context.Background(), 1_000_001, 1_000_002, false, domain.PlanPro, nil)
	require.ErrorIs(t, err, domain.ErrForbidden)
}