123 lines
4.5 KiB
Go
123 lines
4.5 KiB
Go
package repository_test
|
|
|
|
import (
|
|
"context"
|
|
"path/filepath"
|
|
"testing"
|
|
"time"
|
|
|
|
"billing/internal/collection"
|
|
"billing/internal/pricing"
|
|
"billing/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 TestSQLitePricingPersistsCatalogSourceAndManualSaveClearsIt(t *testing.T) {
|
|
store, err := repository.OpenSQLiteUsage(filepath.Join(t.TempDir(), "usage.db"))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer store.Close()
|
|
fetchedAt := time.Date(2026, 8, 15, 10, 0, 0, 0, time.UTC)
|
|
policy := pricing.Policy{
|
|
Model: "deepseek-v4-flash", Base: pricing.Rates{InputMicrosPer1M: 5_000_000, CacheReadMicrosPer1M: 500_000, CacheWriteMicrosPer1M: 6_250_000, OutputMicrosPer1M: 30_000_000},
|
|
FastMultiplier: pricing.Ratio{Numerator: 5, Denominator: 2},
|
|
}
|
|
if err := store.UpsertPriceRecord(context.Background(), repository.PriceRecord{
|
|
Policy: policy,
|
|
Source: repository.PriceSource{Kind: repository.PriceSourceModelsDev, CatalogID: "openai/gpt-5.6-sol", Revision: "revision-one", FetchedAt: fetchedAt},
|
|
}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
records, err := store.ListPriceRecords(context.Background())
|
|
if err != nil || len(records) != 1 {
|
|
t.Fatalf("records=%+v err=%v", records, err)
|
|
}
|
|
if records[0].Source.Kind != repository.PriceSourceModelsDev || records[0].Source.CatalogID != "openai/gpt-5.6-sol" || !records[0].Source.FetchedAt.Equal(fetchedAt) {
|
|
t.Fatalf("catalog source = %+v", records[0].Source)
|
|
}
|
|
policy.Base.InputMicrosPer1M = 7_000_000
|
|
if err := store.UpsertPrice(context.Background(), policy); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
records, err = store.ListPriceRecords(context.Background())
|
|
if err != nil || records[0].Source.Kind != repository.PriceSourceManual || records[0].Source.CatalogID != "" || records[0].Policy.Base.InputMicrosPer1M != 7_000_000 {
|
|
t.Fatalf("manual record=%+v err=%v", records, 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)
|
|
}
|
|
}
|