package repository import ( "context" "errors" "path/filepath" "testing" "time" managedaccess "billing/internal/access" "billing/internal/collection" ) func TestBillingQuotaConcurrencySettlementAndReset(t *testing.T) { ctx := context.Background() store, err := OpenSQLiteUsage(filepath.Join(t.TempDir(), "billing.db")) if err != nil { t.Fatal(err) } defer store.Close() if _, err := store.BootstrapManagedKey(ctx, "default", "000000"); err != nil { t.Fatal(err) } now := time.Now().UTC().Add(time.Second) state, err := store.BillingState(ctx, "key_default", now) if err != nil || state.QuotaMicros != 0 || state.MaxConcurrency != 4 { t.Fatalf("initial billing state=%+v err=%v", state, err) } if _, err := store.AuthorizeBilling(ctx, "key_default", "zero", now); !errors.Is(err, ErrBillingQuotaExhausted) { t.Fatalf("zero quota authorization err=%v", err) } state, err = store.UpdateBilling(ctx, "key_default", managedaccess.BillingSettings{ QuotaMicros: 1_000_000, ResetPeriod: managedaccess.ResetNone, MaxConcurrency: 2, }, now) if err != nil || state.BalanceMicros != 1_000_000 { t.Fatalf("updated billing state=%+v err=%v", state, err) } for _, requestID := range []string{"one", "two"} { if _, err := store.AuthorizeBilling(ctx, "key_default", requestID, now); err != nil { t.Fatalf("authorize %s: %v", requestID, err) } } if _, err := store.AuthorizeBilling(ctx, "key_default", "three", now); !errors.Is(err, ErrBillingConcurrency) { t.Fatalf("third concurrent authorization err=%v", err) } if err := store.CompleteBillingAdmission(ctx, "one", now); err != nil { t.Fatal(err) } if _, err := store.AuthorizeBilling(ctx, "key_default", "three", now); err != nil { t.Fatalf("released slot not reusable: %v", err) } cost := int64(600_000) record := collection.Record{ManagedKeyID: "key_default", ExecutionID: "exec-1", RequestID: "one", RequestedAt: now, Model: "deepseek-v4-flash", CostMicros: &cost} if err := store.Insert(ctx, record); err != nil { t.Fatal(err) } if err := store.Insert(ctx, record); err != nil { t.Fatal(err) } state, _ = store.BillingState(ctx, "key_default", now) if state.SpentMicros != 600_000 || state.LifetimeSpentMicros != 600_000 || state.BalanceMicros != 400_000 { t.Fatalf("duplicate settlement changed balance: %+v", state) } cost = 500_000 record.ExecutionID = "exec-2" record.CostMicros = &cost if err := store.Insert(ctx, record); err != nil { t.Fatal(err) } state, _ = store.BillingState(ctx, "key_default", now) if state.BalanceMicros != -100_000 || state.LifetimeSpentMicros != 1_100_000 { t.Fatalf("soft quota did not preserve negative balance: %+v", state) } if _, err := store.AuthorizeBilling(ctx, "key_default", "blocked", now); !errors.Is(err, ErrBillingQuotaExhausted) { t.Fatalf("negative balance authorization err=%v", err) } state, err = store.ResetBilling(ctx, "key_default", now.Add(time.Minute)) if err != nil || state.SpentMicros != 0 || state.LifetimeSpentMicros != 1_100_000 || state.BalanceMicros != 1_000_000 || state.CycleSequence != 2 { t.Fatalf("manual reset state=%+v err=%v", state, err) } entries, err := store.ListBillingLedger(ctx, "key_default", 0, 50) if err != nil || len(entries) != 4 { t.Fatalf("ledger=%+v err=%v", entries, err) } } func TestBillingLifetimeSpendBackfillsOnceFromCycles(t *testing.T) { ctx := context.Background() path := filepath.Join(t.TempDir(), "billing.db") store, err := OpenSQLiteUsage(path) if err != nil { t.Fatal(err) } if _, err := store.BootstrapManagedKey(ctx, "default", "000000"); err != nil { t.Fatal(err) } cost := int64(375_000) if err := store.Insert(ctx, collection.Record{ManagedKeyID: "key_default", ExecutionID: "historical", RequestedAt: time.Now().UTC(), CostMicros: &cost}); err != nil { t.Fatal(err) } if _, err := store.db.Exec(`UPDATE billing_accounts SET lifetime_spent_micros=0; DELETE FROM cpa_ext_migrations WHERE name='billing_lifetime_spent_v1'`); err != nil { t.Fatal(err) } if err := store.Close(); err != nil { t.Fatal(err) } store, err = OpenSQLiteUsage(path) if err != nil { t.Fatal(err) } defer store.Close() state, err := store.BillingState(ctx, "key_default", time.Now().UTC()) if err != nil || state.LifetimeSpentMicros != cost { t.Fatalf("backfilled state=%+v err=%v", state, err) } } func TestBillingAutomaticMonthlyResetPreservesAnchorDay(t *testing.T) { location := shanghaiLocation() anchor := time.Date(2026, time.January, 31, 10, 0, 0, 0, location) february := advanceReset(anchor, managedaccess.ResetMonthly, 31).In(location) march := advanceReset(february, managedaccess.ResetMonthly, 31).In(location) if february.Day() != 28 || february.Month() != time.February || march.Day() != 31 || march.Month() != time.March { t.Fatalf("monthly sequence: %s then %s", february, march) } }