cpa-plugin/internal/repository/sqlite_access_test.go

package repository_test

import (
	"context"
	"path/filepath"
	"testing"
	"time"

	managedaccess "billing/internal/access"
	"billing/internal/collection"
	"billing/internal/repository"
)

func TestManagedKeysBootstrapLifecycleAndHistoricalUsage(t *testing.T) {
	store, err := repository.OpenSQLiteUsage(filepath.Join(t.TempDir(), "usage.db"))
	if err != nil {
		t.Fatal(err)
	}
	defer store.Close()
	if err := store.Insert(context.Background(), collection.Record{
		APIKey: "000000", Model: "gpt-5", RequestedAt: time.Now(), TotalTokens: 10,
	}); err != nil {
		t.Fatal(err)
	}
	key, err := store.BootstrapManagedKey(context.Background(), "default", "000000")
	if err != nil {
		t.Fatal(err)
	}
	if key.ID != "key_default" || key.Status != managedaccess.StatusActive || !key.AllModels {
		t.Fatalf("bootstrap key = %+v", key)
	}
	records, err := store.ListRecent(context.Background(), 10)
	if err != nil {
		t.Fatal(err)
	}
	if len(records) != 1 || records[0].ManagedKeyID != key.ID || records[0].KeyAlias != "default" {
		t.Fatalf("historical records = %+v", records)
	}
	key.Status = managedaccess.StatusDisabled
	key.AllModels = false
	key.Models = []string{"gpt-5"}
	if err := store.UpdateManagedKey(context.Background(), key); err != nil {
		t.Fatal(err)
	}
	if err := store.ArchiveManagedKey(context.Background(), key.ID); err != nil {
		t.Fatal(err)
	}
	dashboard, err := store.UsageDashboard(context.Background(), time.Now(), 7)
	if err != nil {
		t.Fatal(err)
	}
	if dashboard.Today.Requests != 1 {
		t.Fatalf("archived key usage should remain in dashboard totals: %+v", dashboard)
	}
	for _, user := range dashboard.Users {
		if user.KeyID == key.ID || user.KeyAlias == key.Name {
			t.Fatalf("archived key should not remain in user dashboard: %+v", dashboard)
		}
	}
	archived, err := store.ManagedKeyByID(context.Background(), key.ID)
	if err != nil || archived.Status != managedaccess.StatusArchived {
		t.Fatalf("archived key = %+v, err=%v", archived, err)
	}
	archived.Status = managedaccess.StatusActive
	if err := store.UpdateManagedKey(context.Background(), archived); err == nil {
		t.Fatal("archived key was restored")
	}
}

func TestManagedKeyStatsUseShanghaiDayBoundary(t *testing.T) {
	store, err := repository.OpenSQLiteUsage(filepath.Join(t.TempDir(), "usage.db"))
	if err != nil {
		t.Fatal(err)
	}
	defer store.Close()
	if _, err := store.BootstrapManagedKey(context.Background(), "default", "000000"); err != nil {
		t.Fatal(err)
	}
	cost := int64(25)
	for _, requestedAt := range []time.Time{
		time.Date(2026, 8, 14, 15, 59, 59, 0, time.UTC),
		time.Date(2026, 8, 14, 16, 0, 1, 0, time.UTC),
	} {
		if err := store.Insert(context.Background(), collection.Record{
			ManagedKeyID: "key_default", RequestedAt: requestedAt, TotalTokens: 10, CostMicros: &cost,
		}); err != nil {
			t.Fatal(err)
		}
	}
	shanghai, _ := time.LoadLocation("Asia/Shanghai")
	today := time.Date(2026, 8, 15, 0, 0, 0, 0, shanghai)
	stats, err := store.KeyStats(context.Background(), "key_default", today)
	if err != nil {
		t.Fatal(err)
	}
	if stats.Total.Requests != 2 || stats.Today.Requests != 1 || stats.Today.TotalTokens != 10 || stats.Total.CostMicros != 50 {
		t.Fatalf("stats = %+v", stats)
	}
}

func TestReconcileUpstreamAccountsListsOnlyCurrentSnapshot(t *testing.T) {
	store, err := repository.OpenSQLiteUsage(filepath.Join(t.TempDir(), "usage.db"))
	if err != nil {
		t.Fatal(err)
	}
	defer store.Close()
	ctx := context.Background()
	old := managedaccess.UpstreamAccount{
		CPAAuthID: "codex-user-prolite.json", CPAAuthIndex: "old-index",
		Provider: "codex", DisplayName: "user@example.com", LastSeenAt: time.Now().Add(-time.Hour),
	}
	current := managedaccess.UpstreamAccount{
		CPAAuthID: "codex-user-pro.json", CPAAuthIndex: "current-index",
		Provider: "codex", DisplayName: "user@example.com", LastSeenAt: time.Now(),
	}
	if err := store.ReconcileUpstreamAccounts(ctx, []managedaccess.UpstreamAccount{old, current}); err != nil {
		t.Fatal(err)
	}
	initial, err := store.ListUpstreamAccounts(ctx)
	if err != nil || len(initial) != 2 {
		t.Fatalf("initial accounts = %+v, err=%v", initial, err)
	}
	var oldID string
	for _, account := range initial {
		if account.CPAAuthID == old.CPAAuthID {
			oldID = account.ID
		}
	}
	if err := store.ReconcileUpstreamAccounts(ctx, []managedaccess.UpstreamAccount{current}); err != nil {
		t.Fatal(err)
	}
	listed, err := store.ListUpstreamAccounts(ctx)
	if err != nil || len(listed) != 1 || listed[0].CPAAuthID != current.CPAAuthID {
		t.Fatalf("current accounts = %+v, err=%v", listed, err)
	}
	if historical, err := store.UpstreamAccountByID(ctx, oldID); err != nil || historical.CPAAuthID != old.CPAAuthID {
		t.Fatalf("historical account = %+v, err=%v", historical, err)
	}
}

func TestManagedKeyStatsCountsOrphanUsageAndAliasLifecycleOnce(t *testing.T) {
	store, err := repository.OpenSQLiteUsage(filepath.Join(t.TempDir(), "usage.db"))
	if err != nil {
		t.Fatal(err)
	}
	defer store.Close()
	if _, err := store.BootstrapManagedKey(context.Background(), "default", "000000"); err != nil {
		t.Fatal(err)
	}

	startedAt := time.Date(2026, 8, 14, 17, 41, 30, 0, time.UTC)
	cost := int64(25)
	if err := store.Insert(context.Background(), collection.Record{
		ManagedKeyID: "key_default", RequestedAt: startedAt.Add(4 * time.Millisecond),
		APIKey: "000000", Model: "deepseek-v4-flash", InputTokens: 90,
		OutputTokens: 7, TotalTokens: 97, CostMicros: &cost,
	}); err != nil {
		t.Fatal(err)
	}
	if err := store.UpsertRequest(context.Background(), collection.RequestRecord{
		ManagedKeyID: "key_default", RequestID: "request-alias", RequestedAt: startedAt,
		CompletedAt: startedAt.Add(time.Second), Model: "deepseek-source-a",
		Outcome: "succeeded", StatusCode: 200,
	}); err != nil {
		t.Fatal(err)
	}

	stats, err := store.KeyStats(context.Background(), "key_default", startedAt.Add(-time.Hour))
	if err != nil {
		t.Fatal(err)
	}
	if stats.Total.Requests != 1 || stats.Total.InputTokens != 90 || stats.Total.OutputTokens != 7 ||
		stats.Total.TotalTokens != 97 || stats.Total.CostMicros != 25 {
		t.Fatalf("stats = %+v", stats)
	}
}