146 lines
4.1 KiB
Go
146 lines
4.1 KiB
Go
package repository
|
|
|
|
import (
|
|
"context"
|
|
"path/filepath"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"billing/internal/collection"
|
|
"billing/internal/pricing"
|
|
)
|
|
|
|
func TestSQLiteSeparatesReaderPoolFromWriter(t *testing.T) {
|
|
store, err := OpenSQLiteUsage(filepath.Join(t.TempDir(), "usage.db"))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer store.Close()
|
|
if store.readDB == nil || store.readDB == store.db {
|
|
t.Fatal("reader pool was not initialized separately")
|
|
}
|
|
if got := store.db.Stats().MaxOpenConnections; got != 1 {
|
|
t.Fatalf("writer connections = %d, want 1", got)
|
|
}
|
|
if got := store.readDB.Stats().MaxOpenConnections; got != 4 {
|
|
t.Fatalf("reader connections = %d, want 4", got)
|
|
}
|
|
}
|
|
|
|
func TestSQLiteWriterProgressesDuringReadTransaction(t *testing.T) {
|
|
store, err := OpenSQLiteUsage(filepath.Join(t.TempDir(), "usage.db"))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer store.Close()
|
|
readTx, err := store.readDB.BeginTx(context.Background(), nil)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer readTx.Rollback()
|
|
var count int
|
|
if err := readTx.QueryRow(`SELECT COUNT(*) FROM request_detail_index`).Scan(&count); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
done := make(chan error, 1)
|
|
go func() {
|
|
done <- store.Insert(context.Background(), collection.Record{
|
|
RequestedAt: time.Now().UTC(), Model: "model", TotalTokens: 1,
|
|
})
|
|
}()
|
|
select {
|
|
case err := <-done:
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
case <-time.After(2 * time.Second):
|
|
t.Fatal("writer was blocked by the independent read transaction")
|
|
}
|
|
}
|
|
|
|
func TestBackfillMissingCostsUsesBoundedIndexedBatches(t *testing.T) {
|
|
store, err := OpenSQLiteUsage(filepath.Join(t.TempDir(), "usage.db"))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer store.Close()
|
|
_, err = store.db.Exec(`
|
|
WITH RECURSIVE sequence(value) AS (
|
|
SELECT 1 UNION ALL SELECT value+1 FROM sequence WHERE value<2505
|
|
)
|
|
INSERT INTO usage_records (requested_at, model, failed, input_tokens, output_tokens, total_tokens)
|
|
SELECT printf('2026-08-15T00:00:%02d.000000000Z', value % 60), 'MODEL-BATCH', 0, 1000, 100, 1100
|
|
FROM sequence`)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
policy := pricing.Policy{
|
|
Model: "model-batch",
|
|
Base: pricing.Rates{InputMicrosPer1M: 1_000_000, OutputMicrosPer1M: 2_000_000},
|
|
FastMultiplier: pricing.Ratio{Numerator: 5, Denominator: 2},
|
|
}
|
|
updated, err := store.BackfillMissingCosts(context.Background(), policy)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if updated != 2505 {
|
|
t.Fatalf("updated = %d, want 2505", updated)
|
|
}
|
|
var missing int
|
|
if err := store.readDB.QueryRow(`SELECT COUNT(*) FROM usage_records WHERE cost_micros IS NULL`).Scan(&missing); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if missing != 0 {
|
|
t.Fatalf("missing costs = %d, want 0", missing)
|
|
}
|
|
|
|
plan := explainPlan(t, store, `
|
|
SELECT id FROM usage_records
|
|
WHERE model = ? COLLATE NOCASE AND cost_micros IS NULL AND id > ? AND id <= ?
|
|
ORDER BY id LIMIT ?`, policy.Model, 0, 3000, costBackfillBatchSize)
|
|
if !strings.Contains(plan, "idx_usage_records_missing_cost_model") {
|
|
t.Fatalf("backfill query did not use missing-cost index: %s", plan)
|
|
}
|
|
}
|
|
|
|
func TestUsageDashboardUsesBoundedHistoryIndexes(t *testing.T) {
|
|
store, err := OpenSQLiteUsage(filepath.Join(t.TempDir(), "usage.db"))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer store.Close()
|
|
plan := explainPlan(t, store, usageDashboardUsersSQL, "2026-08-15T00:00:00Z")
|
|
if !strings.Contains(plan, "idx_request_detail_time") {
|
|
t.Fatalf("dashboard today query did not use time index: %s", plan)
|
|
}
|
|
if !strings.Contains(plan, "idx_request_detail_key") {
|
|
t.Fatalf("dashboard latest-use query did not use key index: %s", plan)
|
|
}
|
|
}
|
|
|
|
func explainPlan(t *testing.T, store *SQLiteUsageRepository, statement string, args ...any) string {
|
|
t.Helper()
|
|
rows, err := store.readDB.Query("EXPLAIN QUERY PLAN "+statement, args...)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer rows.Close()
|
|
var plan strings.Builder
|
|
for rows.Next() {
|
|
var id, parent, unused int
|
|
var detail string
|
|
if err := rows.Scan(&id, &parent, &unused, &detail); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
plan.WriteString(detail)
|
|
plan.WriteByte('\n')
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
return plan.String()
|
|
}
|