package repository_test import ( "context" "path/filepath" "testing" "time" "cpa-ext/internal/collection" "cpa-ext/internal/pricing" "cpa-ext/internal/repository" ) func TestSQLitePricingRoundTripAndDelete(t *testing.T) { store, err := repository.OpenSQLiteUsage(filepath.Join(t.TempDir(), "usage.db")) if err != nil { t.Fatal(err) } defer store.Close() policy := pricing.Policy{ Model: "gpt-5.6-sol", Base: pricing.Rates{InputMicrosPer1M: 2_500_000, CacheReadMicrosPer1M: 250_000, CacheWriteMicrosPer1M: 3_125_000, OutputMicrosPer1M: 15_000_000}, LongContext: &pricing.LongContext{ThresholdInputTokens: 272_000, Comparison: "gt", Rates: pricing.Rates{InputMicrosPer1M: 5_000_000, OutputMicrosPer1M: 22_500_000}}, FastPricingEnabled: true, FastMultiplier: pricing.Ratio{Numerator: 5, Denominator: 2}, } if err := store.UpsertPrice(context.Background(), policy); err != nil { t.Fatal(err) } prices, err := store.ListPrices(context.Background()) if err != nil { t.Fatal(err) } if len(prices) != 1 || prices[0].Model != policy.Model || prices[0].LongContext == nil || !prices[0].FastPricingEnabled || prices[0].FastMultiplier != policy.FastMultiplier { t.Fatalf("unexpected prices: %+v", prices) } if err := store.DeletePrice(context.Background(), policy.Model); err != nil { t.Fatal(err) } prices, err = store.ListPrices(context.Background()) if err != nil || len(prices) != 0 { t.Fatalf("prices after delete: %+v, %v", prices, err) } } func TestBackfillMissingCostsDoesNotRewriteExistingCost(t *testing.T) { store, err := repository.OpenSQLiteUsage(filepath.Join(t.TempDir(), "usage.db")) if err != nil { t.Fatal(err) } defer store.Close() missing := collection.Record{RequestedAt: time.Now(), Model: "deepseek-v4-flash", InputTokens: 1_000_000, OutputTokens: 100_000} existingCost := int64(7) priced := collection.Record{RequestedAt: time.Now(), Model: "deepseek-v4-flash", InputTokens: 1_000_000, CostMicros: &existingCost} if err := store.Insert(context.Background(), missing); err != nil { t.Fatal(err) } if err := store.Insert(context.Background(), priced); err != nil { t.Fatal(err) } policy := pricing.Policy{ Model: "deepseek-v4-flash", Base: pricing.Rates{InputMicrosPer1M: 2_500_000, OutputMicrosPer1M: 15_000_000}, FastMultiplier: pricing.Ratio{Numerator: 5, Denominator: 2}, } updated, err := store.BackfillMissingCosts(context.Background(), policy) if err != nil { t.Fatal(err) } if updated != 1 { t.Fatalf("updated = %d, want 1", updated) } records, err := store.ListRecent(context.Background(), 10) if err != nil { t.Fatal(err) } if len(records) != 2 || records[0].CostMicros == nil || records[1].CostMicros == nil { t.Fatalf("unexpected records: %+v", records) } got := map[int64]int{} for _, record := range records { got[*record.CostMicros]++ } if got[4_000_000] != 1 || got[7] != 1 { t.Fatalf("costs = %#v", got) } }