253 lines
8.3 KiB
Go
253 lines
8.3 KiB
Go
package repository_test
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"path/filepath"
|
|
"testing"
|
|
"time"
|
|
|
|
"cpa-ext/internal/collection"
|
|
"cpa-ext/internal/repository"
|
|
)
|
|
|
|
func TestSQLiteUsagePersistsAllRecordsAndLimitsQueries(t *testing.T) {
|
|
databasePath := filepath.Join(t.TempDir(), "usage.db")
|
|
store, err := repository.OpenSQLiteUsage(databasePath)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
requestedAt := time.Date(2026, 8, 14, 12, 0, 0, 0, time.UTC)
|
|
for index := 1; index <= 1002; index++ {
|
|
record := collection.Record{
|
|
RequestedAt: requestedAt.Add(time.Duration(index) * time.Second),
|
|
Model: fmt.Sprintf("model-%d", index),
|
|
TotalTokens: int64(index),
|
|
}
|
|
if err := store.Insert(context.Background(), record); err != nil {
|
|
t.Fatalf("insert record %d: %v", index, err)
|
|
}
|
|
}
|
|
|
|
recent, err := store.ListRecent(context.Background(), 1000)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(recent) != 1000 || recent[0].Model != "model-1002" || recent[999].Model != "model-3" {
|
|
t.Fatalf("unexpected recent records: len=%d first=%q last=%q", len(recent), recent[0].Model, recent[len(recent)-1].Model)
|
|
}
|
|
if err := store.Close(); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
reopened, err := repository.OpenSQLiteUsage(databasePath)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer reopened.Close()
|
|
all, err := reopened.ListRecent(context.Background(), 2000)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(all) != 1002 {
|
|
t.Fatalf("persisted records = %d, want 1002", len(all))
|
|
}
|
|
}
|
|
|
|
func TestSQLiteUsageRoundTripsNullableCostAndDurations(t *testing.T) {
|
|
store, err := repository.OpenSQLiteUsage(filepath.Join(t.TempDir(), "usage.db"))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer store.Close()
|
|
|
|
cost := int64(125_000)
|
|
want := collection.Record{
|
|
RequestedAt: time.Date(2026, 8, 14, 12, 0, 0, 123, time.FixedZone("CST", 8*60*60)),
|
|
APIKey: "key",
|
|
Model: "model",
|
|
Failed: true,
|
|
RequestType: "SSE",
|
|
Endpoint: "POST /v1/responses",
|
|
InputTokens: 10,
|
|
CacheReadTokens: 8,
|
|
TTFT: 250 * time.Millisecond,
|
|
Latency: 2 * time.Second,
|
|
CostMicros: &cost,
|
|
PriceTier: "base",
|
|
PriceMultiplierNumerator: 1,
|
|
PriceMultiplierDenominator: 1,
|
|
ClientIP: "192.0.2.10",
|
|
}
|
|
if err := store.Insert(context.Background(), want); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
records, err := store.ListRecent(context.Background(), 1)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
got := records[0]
|
|
if !got.RequestedAt.Equal(want.RequestedAt) || got.APIKey != want.APIKey || got.RequestType != want.RequestType || got.Endpoint != want.Endpoint || got.ClientIP != want.ClientIP || got.TTFT != want.TTFT || got.Latency != want.Latency || got.CostMicros == nil || *got.CostMicros != cost || got.PriceTier != "base" {
|
|
t.Fatalf("round trip mismatch: %+v", got)
|
|
}
|
|
}
|
|
|
|
func TestSQLiteUsageMergesOrphanUsageWithLifecycleProjection(t *testing.T) {
|
|
for _, outcome := range []struct {
|
|
name string
|
|
value string
|
|
statusCode int
|
|
failed bool
|
|
}{
|
|
{name: "succeeded", value: "succeeded", statusCode: 200},
|
|
{name: "failed", value: "failed", statusCode: 502, failed: true},
|
|
{name: "canceled", value: "canceled", statusCode: 499, failed: true},
|
|
} {
|
|
for _, order := range []string{"usage-first", "lifecycle-first"} {
|
|
t.Run(outcome.name+"/"+order, func(t *testing.T) {
|
|
store, err := repository.OpenSQLiteUsage(filepath.Join(t.TempDir(), "usage.db"))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer store.Close()
|
|
|
|
startedAt := time.Date(2026, 8, 14, 15, 26, 20, 593_000_000, time.UTC)
|
|
cost := int64(12_345)
|
|
usage := collection.Record{
|
|
RequestedAt: startedAt.Add(5 * time.Millisecond), APIKey: "000000", Model: "deepseek-v4-flash",
|
|
ExecutorType: "CodexExecutor", InputTokens: 19_464, OutputTokens: 79, TotalTokens: 19_543,
|
|
TTFT: 266 * time.Millisecond, CostMicros: &cost,
|
|
}
|
|
lifecycle := collection.RequestRecord{
|
|
RequestID: "request-1", TraceID: "trace-1", RequestedAt: startedAt,
|
|
CompletedAt: startedAt.Add(2 * time.Second), Model: "deepseek-v4-flash",
|
|
SourceFormat: "openai-response", Stream: true, Outcome: outcome.value,
|
|
StatusCode: outcome.statusCode, Endpoint: "/v1/responses",
|
|
}
|
|
insertUsage := func() {
|
|
if err := store.Insert(context.Background(), usage); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
insertLifecycle := func() {
|
|
if err := store.UpsertRequest(context.Background(), lifecycle); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
if order == "usage-first" {
|
|
insertUsage()
|
|
insertLifecycle()
|
|
} else {
|
|
insertLifecycle()
|
|
insertUsage()
|
|
}
|
|
|
|
records, err := store.ListRecent(context.Background(), 10)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(records) != 1 {
|
|
t.Fatalf("records = %d, want one logical request: %+v", len(records), records)
|
|
}
|
|
got := records[0]
|
|
if got.RequestID != "request-1" || got.TraceID != "trace-1" || got.APIKey != "000000" ||
|
|
got.RequestType != "SSE" || got.Endpoint != "/v1/responses" || got.Outcome != outcome.value ||
|
|
got.StatusCode != outcome.statusCode || got.Failed != outcome.failed ||
|
|
got.TotalTokens != 19_543 || got.CostMicros == nil || *got.CostMicros != cost {
|
|
t.Fatalf("merged record = %+v", got)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestSQLiteUsageMergesResolvedModelWithManagedKeyAlias(t *testing.T) {
|
|
store, err := repository.OpenSQLiteUsage(filepath.Join(t.TempDir(), "usage.db"))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer store.Close()
|
|
|
|
startedAt := time.Date(2026, 8, 14, 17, 41, 30, 0, time.UTC)
|
|
if err := store.Insert(context.Background(), collection.Record{
|
|
ManagedKeyID: "key-alice", RequestedAt: startedAt.Add(4 * time.Millisecond),
|
|
APIKey: "alice-000000", Model: "deepseek-v4-flash", TotalTokens: 97,
|
|
}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := store.UpsertRequest(context.Background(), collection.RequestRecord{
|
|
ManagedKeyID: "key-alice", RequestID: "request-alias", RequestedAt: startedAt,
|
|
CompletedAt: startedAt.Add(time.Second), Model: "deepseek-source-a", Stream: true,
|
|
Outcome: "succeeded", StatusCode: 200, Endpoint: "/v1/responses",
|
|
}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
records, err := store.ListRecent(context.Background(), 10)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(records) != 1 || records[0].RequestID != "request-alias" || records[0].Model != "deepseek-v4-flash" || records[0].TotalTokens != 97 {
|
|
t.Fatalf("alias records = %+v", records)
|
|
}
|
|
}
|
|
|
|
func TestSQLiteUsageLeavesAmbiguousConcurrentRequestsSeparate(t *testing.T) {
|
|
store, err := repository.OpenSQLiteUsage(filepath.Join(t.TempDir(), "usage.db"))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer store.Close()
|
|
|
|
startedAt := time.Date(2026, 8, 14, 12, 0, 0, 0, time.UTC)
|
|
for index := 0; index < 2; index++ {
|
|
if err := store.Insert(context.Background(), collection.Record{
|
|
RequestedAt: startedAt.Add(5 * time.Millisecond), APIKey: fmt.Sprintf("key-%d", index),
|
|
Model: "deepseek-v4-flash", TotalTokens: int64(100 + index),
|
|
}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := store.UpsertRequest(context.Background(), collection.RequestRecord{
|
|
RequestID: fmt.Sprintf("request-%d", index), RequestedAt: startedAt,
|
|
CompletedAt: startedAt.Add(time.Second), Model: "deepseek-v4-flash", Stream: true,
|
|
Outcome: "succeeded", Endpoint: "/v1/responses",
|
|
}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
|
|
records, err := store.ListRecent(context.Background(), 10)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(records) != 4 {
|
|
t.Fatalf("ambiguous records = %d, want four unmerged facts: %+v", len(records), records)
|
|
}
|
|
}
|
|
|
|
func TestSQLiteUsageKeepsLifecycleWithoutUsageVisible(t *testing.T) {
|
|
store, err := repository.OpenSQLiteUsage(filepath.Join(t.TempDir(), "usage.db"))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer store.Close()
|
|
|
|
startedAt := time.Date(2026, 8, 14, 12, 0, 0, 0, time.UTC)
|
|
if err := store.UpsertRequest(context.Background(), collection.RequestRecord{
|
|
RequestID: "canceled", RequestedAt: startedAt, CompletedAt: startedAt.Add(time.Second),
|
|
Model: "deepseek-v4-flash", Stream: true, Outcome: "canceled", StatusCode: 499,
|
|
Endpoint: "/v1/responses",
|
|
}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
records, err := store.ListRecent(context.Background(), 10)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(records) != 1 || records[0].RequestID != "canceled" || records[0].Outcome != "canceled" {
|
|
t.Fatalf("records = %+v", records)
|
|
}
|
|
}
|