llmgw/internal/store/rollup_test.go

package store

import (
	"context"
	"database/sql"
	"path/filepath"
	"testing"
	"time"
)

// onBoth runs a test on SQLite, and on PostgreSQL when LLMGW_TEST_POSTGRES
// names one: the rollup is written with an upsert, and the two dialects are
// where an upsert goes wrong.
func onBoth(t *testing.T, f func(t *testing.T, s *DB)) {
	t.Run("sqlite", func(t *testing.T) { f(t, openDB(t)) })
	t.Run("postgres", func(t *testing.T) { f(t, openPG(t)) })
}

// slice is a moment a minute into an hour, a few hours back: inside a
// five-minute slice and an hour alike.
func slice() time.Time {
	return time.UnixMilli(tier1h.floor(time.Now().Add(-3 * time.Hour))).Add(time.Minute)
}

func TestTheRollupAddsUpToTheRows(t *testing.T) {
	onBoth(t, func(t *testing.T, s *DB) {
		ctx := context.Background()
		at := slice()
		// Two batches into the same slice, so the second has to add to the
		// row the first made rather than replace it.
		for _, batch := range [][]Usage{
			{
				{At: at, KeyID: "k", UserID: "u", ProjectID: "p", Model: "m", ServedModel: "m",
					PromptTokens: 10, CompletionTokens: 5, CostMicros: 100, Status: 200, LatencyMs: 80, TTFTMs: 40},
				{At: at.Add(time.Second), KeyID: "k", UserID: "u", ProjectID: "p", Model: "m", ServedModel: "m",
					PromptTokens: 1, CompletionTokens: 1, CostMicros: 7, Status: 500, LatencyMs: 3000},
			},
			{
				{At: at.Add(2 * time.Second), KeyID: "k", UserID: "u", ProjectID: "p", Model: "m", ServedModel: "m",
					PromptTokens: 2, CompletionTokens: 3, CostMicros: 50, Status: 200, LatencyMs: 400, TTFTMs: 200},
			},
		} {
			if err := s.WriteUsage(ctx, batch); err != nil {
				t.Fatal(err)
			}
		}

		for _, tr := range tiers {
			var req, errs, prompt, cost, lat, ttftN, last, lat0, lat2, lat5, ttft0, ttft1 int64
			err := s.db.QueryRowContext(ctx, `SELECT requests, errors, prompt_tokens, cost_micros,
			latency_ms_sum, ttft_n, last_ts, lat_0, lat_2, lat_5, ttft_0, ttft_1 FROM `+tr.table).
				Scan(&req, &errs, &prompt, &cost, &lat, &ttftN, &last, &lat0, &lat2, &lat5, &ttft0, &ttft1)
			if err != nil {
				t.Fatalf("%s: %v", tr.table, err)
			}
			if req != 3 || errs != 1 || prompt != 13 || cost != 157 || lat != 3480 {
				t.Errorf("sums: %d requests, %d errors, %d prompt, %d cost, %d ms; want 3, 1, 13, 157, 3480",
					req, errs, prompt, cost, lat)
			}
			// 80 ms in <=100, 400 in <=500, 3000 in <=5000; first tokens 40 and
			// 200, and the failed request had none to count.
			if lat0 != 1 || lat2 != 1 || lat5 != 1 || ttftN != 2 || ttft0 != 1 || ttft1 != 1 {
				t.Errorf("histograms: lat %d/%d/%d, ttft n=%d %d/%d", lat0, lat2, lat5, ttftN, ttft0, ttft1)
			}
			if last != millis(at.Add(2*time.Second)) {
				t.Errorf("last seen %d, want the latest request", last)
			}
		}

		b, err := s.UsageBuckets(ctx, at.Add(-time.Hour), time.Hour)
		if err != nil {
			t.Fatal(err)
		}
		var total Bucket
		for _, x := range b {
			total.Requests += x.Requests
			total.CostMicros += x.CostMicros
		}
		if total.Requests != 3 || total.CostMicros != 157 {
			t.Errorf("the chart counts %d requests and %d spent, want 3 and 157", total.Requests, total.CostMicros)
		}
	})
}

func TestABudgetIsCountedFromTheMomentNotTheSlice(t *testing.T) {
	// Charts start at the slice a window falls in; a budget must not. A
	// request a minute before the period began is not the period's.
	onBoth(t, func(t *testing.T, s *DB) {
		ctx := context.Background()
		at := slice()
		// Before the moment; after it in its own slice; in a later slice of
		// the same hour; and in the next hour — one from each part a budget
		// is read from.
		if err := s.WriteUsage(ctx, []Usage{
			{At: at, KeyID: "k", ProjectID: "p", Model: "m", CostMicros: 100, Status: 200},
			{At: at.Add(2 * time.Minute), KeyID: "k", ProjectID: "p", Model: "m", CostMicros: 10, Status: 200},
			{At: at.Add(10 * time.Minute), KeyID: "k", ProjectID: "p", Model: "m", CostMicros: 1, Status: 200},
			{At: at.Add(70 * time.Minute), KeyID: "k", ProjectID: "p", Model: "m", CostMicros: 1000, Status: 200},
		}); err != nil {
			t.Fatal(err)
		}
		spend, err := s.SpendSince(ctx, at.Add(time.Minute))
		if err != nil {
			t.Fatal(err)
		}
		if spend["k"] != 1011 {
			t.Errorf("spent %d since mid-slice, want 1011", spend["k"])
		}
	})
}

func TestPruningRequestsKeepsEveryTotal(t *testing.T) {
	onBoth(t, func(t *testing.T, s *DB) {
		ctx := context.Background()
		at := slice()
		if err := s.WriteUsage(ctx, []Usage{
			{At: at, KeyID: "k", UserID: "u", ProjectID: "p", Model: "m", ServedModel: "m",
				CostMicros: 100, Status: 200, RequestID: "req-1"},
		}); err != nil {
			t.Fatal(err)
		}
		n, err := s.PruneUsage(ctx, time.Now())
		if err != nil || n != 1 {
			t.Fatalf("pruned %d, %v", n, err)
		}
		if _, ok, _ := s.UsageByRequest(ctx, "req-1"); ok {
			t.Error("a pruned request can still be looked up")
		}
		spend, _ := s.SpendSince(ctx, time.UnixMilli(tier5m.floor(at)))
		users, _ := s.UsageByUser(ctx, at.Add(-time.Hour))
		if spend["k"] != 100 || len(users) != 1 || users[0].Requests != 1 {
			t.Errorf("after pruning: spend %v, people %+v; the totals should not move", spend, users)
		}
	})
}

func TestTheMigrationRollsUpWhatIsAlreadyRecorded(t *testing.T) {
	ctx := context.Background()
	db, err := sql.Open("sqlite", filepath.Join(t.TempDir(), "old.db")+"?_pragma=foreign_keys(1)")
	if err != nil {
		t.Fatal(err)
	}
	defer db.Close()
	db.SetMaxOpenConns(1)
	if _, err := db.ExecContext(ctx,
		`CREATE TABLE schema_migrations (version TEXT PRIMARY KEY, applied_at INTEGER NOT NULL)`); err != nil {
		t.Fatal(err)
	}
	names, _ := migrationNames("sqlite")
	for _, name := range names {
		if name >= "0015" {
			break
		}
		body, _ := migrationFS.ReadFile("migrations/sqlite/" + name)
		if err := applyOne(ctx, db, "sqlite", name, string(body)); err != nil {
			t.Fatalf("%s: %v", name, err)
		}
	}
	for _, ts := range []int64{600_001, 600_002, 900_001} {
		if _, err := db.ExecContext(ctx, `INSERT INTO usage (ts, key_id, project_id, model, served_model,
			prompt_tokens, completion_tokens, cost_micros, status, latency_ms, user_id, ttft_ms)
			VALUES (?, 'k', 'p', 'm', 'm', 1, 1, 5, 200, 700, 'u', 0)`, ts); err != nil {
			t.Fatal(err)
		}
	}
	if err := migrate(ctx, db, "sqlite"); err != nil {
		t.Fatal(err)
	}
	// 600,001 and 600,002 share a five-minute slice and 900,001 starts the
	// next; all three are in the first hour, and the first day.
	for table, want := range map[string]int64{"usage_5m": 2, "usage_1h": 1, "usage_1d": 1} {
		var slices, reqs, cost, lat3, ttftN int64
		if err := db.QueryRowContext(ctx, `SELECT COUNT(*), SUM(requests), SUM(cost_micros), SUM(lat_3), SUM(ttft_n)
			FROM `+table).Scan(&slices, &reqs, &cost, &lat3, &ttftN); err != nil {
			t.Fatal(err)
		}
		if slices != want || reqs != 3 || cost != 15 || lat3 != 3 || ttftN != 0 {
			t.Errorf("%s: %d slices, %d requests, %d spent, %d in <=1000ms, %d first tokens; want %d, 3, 15, 3, 0",
				table, slices, reqs, cost, lat3, ttftN, want)
		}
	}
}

func TestATotalReadsEachTierForItsOwnStretchOnly(t *testing.T) {
	// A total since a moment two days back, a few minutes into a slice, is
	// read from the rows, then five-minute slices, then hours, then days.
	// One request lands in each stretch, and one just before the moment;
	// counting any of them twice, or the early one at all, shows here.
	onBoth(t, func(t *testing.T, s *DB) {
		ctx := context.Background()
		day := time.UnixMilli(tier1d.floor(time.Now().Add(-48 * time.Hour)))
		since := day.Add(22*time.Hour + 52*time.Minute + 30*time.Second)
		var batch []Usage
		for i, at := range []time.Time{
			since.Add(-time.Second),     // before: not counted
			since.Add(time.Minute),      // the rows
			since.Add(5 * time.Minute),  // five-minute slices
			since.Add(20 * time.Minute), // hours
			since.Add(2 * time.Hour),    // days
		} {
			batch = append(batch, Usage{At: at, KeyID: "k", UserID: "u", ProjectID: "p", Model: "m",
				ServedModel: "m", CostMicros: int64(1) << i, Status: 200})
		}
		if err := s.WriteUsage(ctx, batch); err != nil {
			t.Fatal(err)
		}
		users, err := s.UsageByUser(ctx, since)
		if err != nil {
			t.Fatal(err)
		}
		spend, err := s.SpendSince(ctx, since)
		if err != nil {
			t.Fatal(err)
		}
		if len(users) != 1 || users[0].Requests != 4 || users[0].CostMicros != 2+4+8+16 || spend["k"] != 30 {
			t.Errorf("people %+v, spend %d; want 4 requests and 30 spent", users, spend["k"])
		}
	})
}