cpa-plugin/internal/repository/sqlite_billing_test.go

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