| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367 |
- //go:build integration
- package repository
- import (
- "context"
- "fmt"
- "strings"
- "testing"
- "time"
- "github.com/google/uuid"
- "github.com/stretchr/testify/require"
- "github.com/Wei-Shaw/sub2api/internal/service"
- )
- func TestUsageBillingRepositoryApply_DeduplicatesBalanceBilling(t *testing.T) {
- ctx := context.Background()
- client := testEntClient(t)
- repo := NewUsageBillingRepository(client, integrationDB)
- user := mustCreateUser(t, client, &service.User{
- Email: fmt.Sprintf("usage-billing-user-%d@example.com", time.Now().UnixNano()),
- PasswordHash: "hash",
- Balance: 100,
- })
- apiKey := mustCreateApiKey(t, client, &service.APIKey{
- UserID: user.ID,
- Key: "sk-usage-billing-" + uuid.NewString(),
- Name: "billing",
- Quota: 1,
- })
- account := mustCreateAccount(t, client, &service.Account{
- Name: "usage-billing-account-" + uuid.NewString(),
- Type: service.AccountTypeAPIKey,
- })
- requestID := uuid.NewString()
- cmd := &service.UsageBillingCommand{
- RequestID: requestID,
- APIKeyID: apiKey.ID,
- UserID: user.ID,
- AccountID: account.ID,
- AccountType: service.AccountTypeAPIKey,
- BalanceCost: 1.25,
- APIKeyQuotaCost: 1.25,
- APIKeyRateLimitCost: 1.25,
- }
- result1, err := repo.Apply(ctx, cmd)
- require.NoError(t, err)
- require.NotNil(t, result1)
- require.True(t, result1.Applied)
- require.True(t, result1.APIKeyQuotaExhausted)
- result2, err := repo.Apply(ctx, cmd)
- require.NoError(t, err)
- require.NotNil(t, result2)
- require.False(t, result2.Applied)
- var balance float64
- require.NoError(t, integrationDB.QueryRowContext(ctx, "SELECT balance FROM users WHERE id = $1", user.ID).Scan(&balance))
- require.InDelta(t, 98.75, balance, 0.000001)
- var quotaUsed float64
- require.NoError(t, integrationDB.QueryRowContext(ctx, "SELECT quota_used FROM api_keys WHERE id = $1", apiKey.ID).Scan("aUsed))
- require.InDelta(t, 1.25, quotaUsed, 0.000001)
- var usage5h float64
- require.NoError(t, integrationDB.QueryRowContext(ctx, "SELECT usage_5h FROM api_keys WHERE id = $1", apiKey.ID).Scan(&usage5h))
- require.InDelta(t, 1.25, usage5h, 0.000001)
- var status string
- require.NoError(t, integrationDB.QueryRowContext(ctx, "SELECT status FROM api_keys WHERE id = $1", apiKey.ID).Scan(&status))
- require.Equal(t, service.StatusAPIKeyQuotaExhausted, status)
- var dedupCount int
- require.NoError(t, integrationDB.QueryRowContext(ctx, "SELECT COUNT(*) FROM usage_billing_dedup WHERE request_id = $1 AND api_key_id = $2", requestID, apiKey.ID).Scan(&dedupCount))
- require.Equal(t, 1, dedupCount)
- }
- func TestUsageBillingRepositoryApply_DeduplicatesSubscriptionBilling(t *testing.T) {
- ctx := context.Background()
- client := testEntClient(t)
- repo := NewUsageBillingRepository(client, integrationDB)
- user := mustCreateUser(t, client, &service.User{
- Email: fmt.Sprintf("usage-billing-sub-user-%d@example.com", time.Now().UnixNano()),
- PasswordHash: "hash",
- })
- group := mustCreateGroup(t, client, &service.Group{
- Name: "usage-billing-group-" + uuid.NewString(),
- Platform: service.PlatformAnthropic,
- SubscriptionType: service.SubscriptionTypeSubscription,
- })
- apiKey := mustCreateApiKey(t, client, &service.APIKey{
- UserID: user.ID,
- GroupID: &group.ID,
- Key: "sk-usage-billing-sub-" + uuid.NewString(),
- Name: "billing-sub",
- })
- subscription := mustCreateSubscription(t, client, &service.UserSubscription{
- UserID: user.ID,
- GroupID: group.ID,
- })
- requestID := uuid.NewString()
- cmd := &service.UsageBillingCommand{
- RequestID: requestID,
- APIKeyID: apiKey.ID,
- UserID: user.ID,
- AccountID: 0,
- SubscriptionID: &subscription.ID,
- SubscriptionCost: 2.5,
- }
- result1, err := repo.Apply(ctx, cmd)
- require.NoError(t, err)
- require.True(t, result1.Applied)
- result2, err := repo.Apply(ctx, cmd)
- require.NoError(t, err)
- require.False(t, result2.Applied)
- var dailyUsage float64
- require.NoError(t, integrationDB.QueryRowContext(ctx, "SELECT daily_usage_usd FROM user_subscriptions WHERE id = $1", subscription.ID).Scan(&dailyUsage))
- require.InDelta(t, 2.5, dailyUsage, 0.000001)
- }
- func TestUsageBillingRepositoryApply_RequestFingerprintConflict(t *testing.T) {
- ctx := context.Background()
- client := testEntClient(t)
- repo := NewUsageBillingRepository(client, integrationDB)
- user := mustCreateUser(t, client, &service.User{
- Email: fmt.Sprintf("usage-billing-conflict-user-%d@example.com", time.Now().UnixNano()),
- PasswordHash: "hash",
- Balance: 100,
- })
- apiKey := mustCreateApiKey(t, client, &service.APIKey{
- UserID: user.ID,
- Key: "sk-usage-billing-conflict-" + uuid.NewString(),
- Name: "billing-conflict",
- })
- requestID := uuid.NewString()
- _, err := repo.Apply(ctx, &service.UsageBillingCommand{
- RequestID: requestID,
- APIKeyID: apiKey.ID,
- UserID: user.ID,
- BalanceCost: 1.25,
- })
- require.NoError(t, err)
- _, err = repo.Apply(ctx, &service.UsageBillingCommand{
- RequestID: requestID,
- APIKeyID: apiKey.ID,
- UserID: user.ID,
- BalanceCost: 2.50,
- })
- require.ErrorIs(t, err, service.ErrUsageBillingRequestConflict)
- }
- func TestUsageBillingRepositoryApply_UpdatesAccountQuota(t *testing.T) {
- ctx := context.Background()
- client := testEntClient(t)
- repo := NewUsageBillingRepository(client, integrationDB)
- user := mustCreateUser(t, client, &service.User{
- Email: fmt.Sprintf("usage-billing-account-user-%d@example.com", time.Now().UnixNano()),
- PasswordHash: "hash",
- })
- apiKey := mustCreateApiKey(t, client, &service.APIKey{
- UserID: user.ID,
- Key: "sk-usage-billing-account-" + uuid.NewString(),
- Name: "billing-account",
- })
- account := mustCreateAccount(t, client, &service.Account{
- Name: "usage-billing-account-quota-" + uuid.NewString(),
- Type: service.AccountTypeAPIKey,
- Extra: map[string]any{
- "quota_limit": 100.0,
- },
- })
- _, err := repo.Apply(ctx, &service.UsageBillingCommand{
- RequestID: uuid.NewString(),
- APIKeyID: apiKey.ID,
- UserID: user.ID,
- AccountID: account.ID,
- AccountType: service.AccountTypeAPIKey,
- AccountQuotaCost: 3.5,
- })
- require.NoError(t, err)
- var quotaUsed float64
- require.NoError(t, integrationDB.QueryRowContext(ctx, "SELECT COALESCE((extra->>'quota_used')::numeric, 0) FROM accounts WHERE id = $1", account.ID).Scan("aUsed))
- require.InDelta(t, 3.5, quotaUsed, 0.000001)
- }
- func TestUsageBillingRepositoryApply_EnqueuesSchedulerOutboxOnQuotaCrossing(t *testing.T) {
- ctx := context.Background()
- client := testEntClient(t)
- repo := NewUsageBillingRepository(client, integrationDB)
- newFixture := func(t *testing.T, extra map[string]any) (int64, int64) {
- t.Helper()
- user := mustCreateUser(t, client, &service.User{
- Email: fmt.Sprintf("usage-billing-outbox-user-%d-%s@example.com", time.Now().UnixNano(), uuid.NewString()),
- PasswordHash: "hash",
- })
- apiKey := mustCreateApiKey(t, client, &service.APIKey{
- UserID: user.ID,
- Key: "sk-usage-billing-outbox-" + uuid.NewString(),
- Name: "billing-outbox",
- })
- account := mustCreateAccount(t, client, &service.Account{
- Name: "usage-billing-outbox-" + uuid.NewString(),
- Type: service.AccountTypeAPIKey,
- Extra: extra,
- })
- return apiKey.ID, account.ID
- }
- outboxCountFor := func(t *testing.T, accountID int64) int {
- t.Helper()
- var count int
- require.NoError(t, integrationDB.QueryRowContext(ctx,
- "SELECT COUNT(*) FROM scheduler_outbox WHERE event_type = $1 AND account_id = $2",
- service.SchedulerOutboxEventAccountChanged, accountID,
- ).Scan(&count))
- return count
- }
- t.Run("daily_first_crossing_enqueues", func(t *testing.T) {
- apiKeyID, accountID := newFixture(t, map[string]any{
- "quota_daily_limit": 10.0,
- })
- // 第一次低于日限额:不应入队 outbox
- _, err := repo.Apply(ctx, &service.UsageBillingCommand{
- RequestID: uuid.NewString(),
- APIKeyID: apiKeyID,
- AccountID: accountID,
- AccountType: service.AccountTypeAPIKey,
- AccountQuotaCost: 4,
- })
- require.NoError(t, err)
- require.Equal(t, 0, outboxCountFor(t, accountID), "below limit should not enqueue")
- // 第二次跨越日限额:应入队一次 outbox
- _, err = repo.Apply(ctx, &service.UsageBillingCommand{
- RequestID: uuid.NewString(),
- APIKeyID: apiKeyID,
- AccountID: accountID,
- AccountType: service.AccountTypeAPIKey,
- AccountQuotaCost: 8,
- })
- require.NoError(t, err)
- require.Equal(t, 1, outboxCountFor(t, accountID), "crossing daily limit should enqueue once")
- // 再次递增(已超):不应重复入队
- _, err = repo.Apply(ctx, &service.UsageBillingCommand{
- RequestID: uuid.NewString(),
- APIKeyID: apiKeyID,
- AccountID: accountID,
- AccountType: service.AccountTypeAPIKey,
- AccountQuotaCost: 2,
- })
- require.NoError(t, err)
- require.Equal(t, 1, outboxCountFor(t, accountID), "subsequent increments beyond limit should not re-enqueue")
- })
- t.Run("weekly_first_crossing_enqueues", func(t *testing.T) {
- apiKeyID, accountID := newFixture(t, map[string]any{
- "quota_weekly_limit": 10.0,
- })
- _, err := repo.Apply(ctx, &service.UsageBillingCommand{
- RequestID: uuid.NewString(),
- APIKeyID: apiKeyID,
- AccountID: accountID,
- AccountType: service.AccountTypeAPIKey,
- AccountQuotaCost: 15, // 单次即跨越
- })
- require.NoError(t, err)
- require.Equal(t, 1, outboxCountFor(t, accountID), "single-shot crossing weekly limit should enqueue once")
- })
- }
- func TestDashboardAggregationRepositoryCleanupUsageBillingDedup_BatchDeletesOldRows(t *testing.T) {
- ctx := context.Background()
- repo := newDashboardAggregationRepositoryWithSQL(integrationDB)
- oldRequestID := "dedup-old-" + uuid.NewString()
- newRequestID := "dedup-new-" + uuid.NewString()
- oldCreatedAt := time.Now().UTC().AddDate(0, 0, -400)
- newCreatedAt := time.Now().UTC().Add(-time.Hour)
- _, err := integrationDB.ExecContext(ctx, `
- INSERT INTO usage_billing_dedup (request_id, api_key_id, request_fingerprint, created_at)
- VALUES ($1, 1, $2, $3), ($4, 1, $5, $6)
- `,
- oldRequestID, strings.Repeat("a", 64), oldCreatedAt,
- newRequestID, strings.Repeat("b", 64), newCreatedAt,
- )
- require.NoError(t, err)
- require.NoError(t, repo.CleanupUsageBillingDedup(ctx, time.Now().UTC().AddDate(0, 0, -365)))
- var oldCount int
- require.NoError(t, integrationDB.QueryRowContext(ctx, "SELECT COUNT(*) FROM usage_billing_dedup WHERE request_id = $1", oldRequestID).Scan(&oldCount))
- require.Equal(t, 0, oldCount)
- var newCount int
- require.NoError(t, integrationDB.QueryRowContext(ctx, "SELECT COUNT(*) FROM usage_billing_dedup WHERE request_id = $1", newRequestID).Scan(&newCount))
- require.Equal(t, 1, newCount)
- var archivedCount int
- require.NoError(t, integrationDB.QueryRowContext(ctx, "SELECT COUNT(*) FROM usage_billing_dedup_archive WHERE request_id = $1", oldRequestID).Scan(&archivedCount))
- require.Equal(t, 1, archivedCount)
- }
- func TestUsageBillingRepositoryApply_DeduplicatesAgainstArchivedKey(t *testing.T) {
- ctx := context.Background()
- client := testEntClient(t)
- repo := NewUsageBillingRepository(client, integrationDB)
- aggRepo := newDashboardAggregationRepositoryWithSQL(integrationDB)
- user := mustCreateUser(t, client, &service.User{
- Email: fmt.Sprintf("usage-billing-archive-user-%d@example.com", time.Now().UnixNano()),
- PasswordHash: "hash",
- Balance: 100,
- })
- apiKey := mustCreateApiKey(t, client, &service.APIKey{
- UserID: user.ID,
- Key: "sk-usage-billing-archive-" + uuid.NewString(),
- Name: "billing-archive",
- })
- requestID := uuid.NewString()
- cmd := &service.UsageBillingCommand{
- RequestID: requestID,
- APIKeyID: apiKey.ID,
- UserID: user.ID,
- BalanceCost: 1.25,
- }
- result1, err := repo.Apply(ctx, cmd)
- require.NoError(t, err)
- require.True(t, result1.Applied)
- _, err = integrationDB.ExecContext(ctx, `
- UPDATE usage_billing_dedup
- SET created_at = $1
- WHERE request_id = $2 AND api_key_id = $3
- `, time.Now().UTC().AddDate(0, 0, -400), requestID, apiKey.ID)
- require.NoError(t, err)
- require.NoError(t, aggRepo.CleanupUsageBillingDedup(ctx, time.Now().UTC().AddDate(0, 0, -365)))
- result2, err := repo.Apply(ctx, cmd)
- require.NoError(t, err)
- require.False(t, result2.Applied)
- var balance float64
- require.NoError(t, integrationDB.QueryRowContext(ctx, "SELECT balance FROM users WHERE id = $1", user.ID).Scan(&balance))
- require.InDelta(t, 98.75, balance, 0.000001)
- }
|