llmgw/internal/store/store_test.go

package store

import (
	"context"
	"errors"
	"os"
	"path/filepath"
	"slices"
	"strings"
	"testing"
	"time"
)

func openDB(t *testing.T) *DB {
	t.Helper()
	s, err := OpenSQLite(context.Background(), filepath.Join(t.TempDir(), "llmgw.db"))
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = s.Close() })
	return s
}

// tree builds org → team → project and returns the project id.
func tree(t *testing.T, s *DB) string {
	t.Helper()
	ctx := context.Background()
	if err := s.PutOrg(ctx, Org{ID: "org1", Name: "Acme"}); err != nil {
		t.Fatal(err)
	}
	if err := s.PutTeam(ctx, Team{ID: "team1", OrgID: "org1", Name: "Research"}); err != nil {
		t.Fatal(err)
	}
	if err := s.PutProject(ctx, Project{ID: "proj1", TeamID: "team1", Name: "Chat"}); err != nil {
		t.Fatal(err)
	}
	return "proj1"
}

func mustKey(t *testing.T, s *DB, id, project string) (string, Key) {
	t.Helper()
	secret, k, err := NewKey(id, id, project)
	if err != nil {
		t.Fatal(err)
	}
	k.CreatedAt = time.Now()
	if err := s.CreateKey(context.Background(), k); err != nil {
		t.Fatal(err)
	}
	return secret, k
}

func intp(v int) *int     { return &v }
func i64p(v int64) *int64 { return &v }

func TestMigrationsAreIdempotent(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "llmgw.db")

	s, err := OpenSQLite(context.Background(), path)
	if err != nil {
		t.Fatal(err)
	}
	tree(t, s)
	_ = s.Close()

	// Reopening must not try to create what is already there, and must not lose
	// what was written.
	s2, err := OpenSQLite(context.Background(), path)
	if err != nil {
		t.Fatalf("reopening ran the migrations again: %v", err)
	}
	defer s2.Close()

	orgs, err := s2.Orgs(context.Background())
	if err != nil || len(orgs) != 1 {
		t.Fatalf("orgs after reopen: %v %v", orgs, err)
	}
}

func TestTheKeyItselfIsNeverStored(t *testing.T) {
	// The whole point of hashing: a database leak hands over no working
	// credential.
	s := openDB(t)
	project := tree(t, s)
	secret, _ := mustKey(t, s, "key1", project)

	rows, err := s.db.QueryContext(context.Background(), `SELECT id, key_hash, key_prefix, name FROM api_keys`)
	if err != nil {
		t.Fatal(err)
	}
	defer rows.Close()
	for rows.Next() {
		var id, prefix, name string
		var hash []byte
		if err := rows.Scan(&id, &hash, &prefix, &name); err != nil {
			t.Fatal(err)
		}
		for _, field := range []string{id, prefix, name, string(hash)} {
			if strings.Contains(field, strings.TrimPrefix(secret, KeyPrefix)) {
				t.Fatal("the key itself reached the database")
			}
		}
		if len(hash) != 32 {
			t.Errorf("hash is %d bytes", len(hash))
		}
	}
}

func TestAuthenticationFindsAKeyByWhatWasPresented(t *testing.T) {
	s := openDB(t)
	project := tree(t, s)
	secret, created := mustKey(t, s, "key1", project)

	got, err := s.KeyByHash(context.Background(), HashKey(secret))
	if err != nil {
		t.Fatal(err)
	}
	if got.ID != created.ID || got.ProjectID != project {
		t.Errorf("found %+v", got)
	}
	if !got.Usable(time.Now()) {
		t.Error("a fresh key is not usable")
	}

	if _, err := s.KeyByHash(context.Background(), HashKey(secret+"x")); !errors.Is(err, ErrNotFound) {
		t.Errorf("a wrong key was accepted: %v", err)
	}
}

func TestRevocationKeepsTheAuditTrail(t *testing.T) {
	// A revoked key must stop working and must still exist: usage rows
	// reference it, and an audit trail that loses its subject is not one.
	s := openDB(t)
	ctx := context.Background()
	project := tree(t, s)
	secret, k := mustKey(t, s, "key1", project)

	if err := s.RevokeKey(ctx, k.ID, time.Now()); err != nil {
		t.Fatal(err)
	}
	got, err := s.KeyByHash(ctx, HashKey(secret))
	if err != nil {
		t.Fatalf("a revoked key vanished from the database: %v", err)
	}
	if got.Usable(time.Now()) {
		t.Error("a revoked key still authenticates")
	}

	if err := s.RevokeKey(ctx, "nope", time.Now()); !errors.Is(err, ErrNotFound) {
		t.Errorf("revoking a key that does not exist: %v", err)
	}
}

func TestExpiryIsCheckedAgainstTheClockNotTheDatabase(t *testing.T) {
	s := openDB(t)
	project := tree(t, s)
	secret, k := mustKey(t, s, "key1", project)

	k.ExpiresAt = time.Now().Add(-time.Hour)
	if err := s.CreateKey(context.Background(), Key{
		ID: "key2", Hash: HashKey("other"), Prefix: "x", Name: "n",
		ProjectID: project, CreatedAt: time.Now(), ExpiresAt: k.ExpiresAt,
	}); err != nil {
		t.Fatal(err)
	}

	expired, err := s.KeyByHash(context.Background(), HashKey("other"))
	if err != nil {
		t.Fatal(err)
	}
	if expired.Usable(time.Now()) {
		t.Error("an expired key is usable")
	}
	fresh, _ := s.KeyByHash(context.Background(), HashKey(secret))
	if !fresh.Usable(time.Now()) {
		t.Error("an unexpired key is not usable")
	}
}

func TestOrphansAreRefused(t *testing.T) {
	// SQLite leaves foreign keys off by default, which would let the tenancy
	// tree grow keys belonging to no project.
	s := openDB(t)
	_, k, _ := NewKey("key1", "n", "no-such-project")
	k.CreatedAt = time.Now()
	if err := s.CreateKey(context.Background(), k); err == nil {
		t.Fatal("a key was created under a project that does not exist")
	}
}

func TestLimitsResolveToTheMostRestrictive(t *testing.T) {
	s := openDB(t)
	ctx := context.Background()
	project := tree(t, s)
	_, k := mustKey(t, s, "key1", project)

	// The organization is generous, the team less so, the key least. Whichever
	// is tightest has to win at every field independently.
	must(t, s.PutLimits(ctx, Limits{Scope: ScopeOrg, ScopeID: "org1", RPM: intp(1000), TPM: intp(100000)}))
	must(t, s.PutLimits(ctx, Limits{Scope: ScopeTeam, ScopeID: "team1", RPM: intp(100)}))
	must(t, s.PutLimits(ctx, Limits{Scope: ScopeKey, ScopeID: k.ID, TPM: intp(500)}))

	e, err := s.Resolve(ctx, k.ID)
	if err != nil {
		t.Fatal(err)
	}
	if e.RPM == nil || *e.RPM != 100 {
		t.Errorf("rpm=%v, the team's limit should win", e.RPM)
	}
	if e.TPM == nil || *e.TPM != 500 {
		t.Errorf("tpm=%v, the key's limit should win", e.TPM)
	}
	if e.OrgID != "org1" || e.TeamID != "team1" || e.ProjectID != project {
		t.Errorf("chain resolved to %+v", e)
	}
}

func TestANarrowerScopeCannotRaiseALimit(t *testing.T) {
	// The only composition rule under which a budget means anything.
	s := openDB(t)
	ctx := context.Background()
	project := tree(t, s)
	_, k := mustKey(t, s, "key1", project)

	must(t, s.PutLimits(ctx, Limits{Scope: ScopeOrg, ScopeID: "org1", RPM: intp(10),
		Budget: i64p(1_000_000), BudgetPeriod: PeriodMonth}))
	must(t, s.PutLimits(ctx, Limits{Scope: ScopeKey, ScopeID: k.ID, RPM: intp(9999),
		Budget: i64p(999_000_000), BudgetPeriod: PeriodDay}))

	e, err := s.Resolve(ctx, k.ID)
	if err != nil {
		t.Fatal(err)
	}
	if *e.RPM != 10 {
		t.Errorf("a key raised its organization's rate limit to %d", *e.RPM)
	}
	if *e.Budget != 1_000_000 || e.BudgetPeriod != PeriodMonth {
		t.Errorf("a key raised its organization's budget to %d/%s", *e.Budget, e.BudgetPeriod)
	}
}

func TestAllowlistsIntersectRatherThanReplace(t *testing.T) {
	s := openDB(t)
	ctx := context.Background()
	project := tree(t, s)
	_, k := mustKey(t, s, "key1", project)

	must(t, s.PutLimits(ctx, Limits{Scope: ScopeTeam, ScopeID: "team1",
		AllowedModels: []string{"qwen", "llama"}}))
	must(t, s.PutLimits(ctx, Limits{Scope: ScopeKey, ScopeID: k.ID,
		AllowedModels: []string{"llama", "gpt-4"}}))

	e, err := s.Resolve(ctx, k.ID)
	if err != nil {
		t.Fatal(err)
	}
	if !slices.Equal(e.AllowedModels, []string{"llama"}) {
		t.Fatalf("allowed=%v; a key must not add a model its team withheld", e.AllowedModels)
	}
	if e.Allows("gpt-4") {
		t.Error("a model the team never allowed is served")
	}
	if !e.Allows("llama") {
		t.Error("the intersection is not served")
	}
}

func TestNoAllowlistMeansEveryModel(t *testing.T) {
	s := openDB(t)
	project := tree(t, s)
	_, k := mustKey(t, s, "key1", project)

	e, err := s.Resolve(context.Background(), k.ID)
	if err != nil {
		t.Fatal(err)
	}
	if e.AllowedModels != nil {
		t.Errorf("allowed=%v", e.AllowedModels)
	}
	if !e.Allows("anything") {
		t.Error("an unrestricted key was refused a model")
	}
}

func TestAnEmptyAllowlistForbidsEverything(t *testing.T) {
	// Empty is not the same as absent, and the store must not conflate them:
	// one is "no restriction", the other is "nothing".
	s := openDB(t)
	ctx := context.Background()
	project := tree(t, s)
	_, k := mustKey(t, s, "key1", project)

	must(t, s.PutLimits(ctx, Limits{Scope: ScopeKey, ScopeID: k.ID, AllowedModels: []string{}}))

	e, err := s.Resolve(ctx, k.ID)
	if err != nil {
		t.Fatal(err)
	}
	if e.AllowedModels == nil {
		t.Fatal("an empty allowlist was read back as no allowlist")
	}
	if e.Allows("qwen") {
		t.Error("an empty allowlist served a model")
	}
}

func TestBudgetSurvivesARestart(t *testing.T) {
	// A counter that lives only in memory resets on restart, which would make a
	// budget a suggestion.
	dir := t.TempDir()
	path := filepath.Join(dir, "llmgw.db")
	ctx := context.Background()

	s, err := OpenSQLite(ctx, path)
	if err != nil {
		t.Fatal(err)
	}
	project := tree(t, s)
	_, k := mustKey(t, s, "key1", project)

	now := time.Now()
	must(t, s.WriteUsage(ctx, []Usage{
		{At: now, KeyID: k.ID, ProjectID: project, Model: "qwen", ServedModel: "qwen",
			PromptTokens: 100, CompletionTokens: 50, CostMicros: 1500, Status: 200},
		{At: now, KeyID: k.ID, ProjectID: project, Model: "qwen", ServedModel: "qwen",
			PromptTokens: 200, CompletionTokens: 80, CostMicros: 2800, Status: 200},
		// Last period: must not count towards this one.
		{At: now.Add(-48 * time.Hour), KeyID: k.ID, ProjectID: project, Model: "qwen",
			ServedModel: "qwen", CostMicros: 999999, Status: 200},
	}))
	_ = s.Close()

	s2, err := OpenSQLite(ctx, path)
	if err != nil {
		t.Fatal(err)
	}
	defer s2.Close()

	spend, err := s2.SpendSince(ctx, now.Add(-time.Hour))
	if err != nil {
		t.Fatal(err)
	}
	if spend[k.ID] != 4300 {
		t.Errorf("spend after restart = %d, want 4300", spend[k.ID])
	}
}

func TestLastUsedIsOnlyEverMovedForward(t *testing.T) {
	// The batch writer sees keys out of order; an older timestamp must not
	// rewind the newer one.
	s := openDB(t)
	ctx := context.Background()
	project := tree(t, s)
	_, k := mustKey(t, s, "key1", project)

	newer := time.Now().Truncate(time.Millisecond)
	older := newer.Add(-time.Hour)

	must(t, s.TouchKeys(ctx, map[string]time.Time{k.ID: newer}))
	must(t, s.TouchKeys(ctx, map[string]time.Time{k.ID: older}))

	got, err := s.KeyByHash(ctx, k.Hash)
	if err != nil {
		t.Fatal(err)
	}
	if !got.LastUsedAt.Equal(newer.UTC()) {
		t.Errorf("last used = %v, want %v", got.LastUsedAt, newer.UTC())
	}
}

func TestUpsertsAreIdempotent(t *testing.T) {
	// Applying a configuration twice must not double the tree or fail.
	s := openDB(t)
	ctx := context.Background()
	tree(t, s)
	tree(t, s)

	orgs, _ := s.Orgs(ctx)
	teams, _ := s.Teams(ctx, "org1")
	projects, _ := s.Projects(ctx, "team1")
	if len(orgs) != 1 || len(teams) != 1 || len(projects) != 1 {
		t.Errorf("tree duplicated: %d orgs, %d teams, %d projects", len(orgs), len(teams), len(projects))
	}

	must(t, s.PutOrg(ctx, Org{ID: "org1", Name: "Renamed"}))
	orgs, _ = s.Orgs(ctx)
	if orgs[0].Name != "Renamed" {
		t.Errorf("upsert did not update: %q", orgs[0].Name)
	}
}

func TestResolveRejectsAnUnknownKey(t *testing.T) {
	s := openDB(t)
	tree(t, s)
	if _, err := s.Resolve(context.Background(), "nope"); !errors.Is(err, ErrNotFound) {
		t.Errorf("resolving an unknown key: %v", err)
	}
}

func must(t *testing.T, err error) {
	t.Helper()
	if err != nil {
		t.Fatal(err)
	}
}

func TestTheStoreIsNotWorldReadable(t *testing.T) {
	// SQLite creates the file with the process umask, which on a stock system
	// leaves it readable by every account on the machine. What is in here is
	// not a working credential, but it is the whole tenancy tree and every
	// request anyone made.
	dir := t.TempDir()
	path := filepath.Join(dir, "llmgw.db")
	s, err := OpenSQLite(context.Background(), path)
	if err != nil {
		t.Fatal(err)
	}
	defer s.Close()

	// Force the WAL files into existence.
	tree(t, s)

	for _, name := range []string{"llmgw.db", "llmgw.db-wal", "llmgw.db-shm"} {
		info, err := os.Stat(filepath.Join(dir, name))
		if os.IsNotExist(err) {
			continue
		}
		if err != nil {
			t.Fatal(err)
		}
		if mode := info.Mode().Perm(); mode&0o077 != 0 {
			t.Errorf("%s is mode %o", name, mode)
		}
	}
}

func TestDefaultLimitsApplyWhereNoScopeSetsOne(t *testing.T) {
	// Unlimited was the default, and a project created without limits is the
	// easy path — the one this cluster's "default" project took. The first
	// sign of that is the bill or the outage.
	s := openDB(t)
	proj := tree(t, s)
	_, k := mustKey(t, s, "k1", proj)

	s.SetDefaultLimits(intp(60), intp(90_000))

	e, err := s.Resolve(context.Background(), k.ID)
	if err != nil {
		t.Fatal(err)
	}
	if e.RPM == nil || *e.RPM != 60 {
		t.Errorf("RPM = %v, want the default", e.RPM)
	}
	if e.TPM == nil || *e.TPM != 90_000 {
		t.Errorf("TPM = %v, want the default", e.TPM)
	}
}

func TestAnExplicitLimitBeatsTheDefaultInBothDirections(t *testing.T) {
	// A default is a fallback, not a ceiling. Folding it in as the widest scope
	// and tightening would silently cap a project deliberately given more.
	s := openDB(t)
	proj := tree(t, s)
	_, k := mustKey(t, s, "k1", proj)

	ctx := context.Background()
	if err := s.PutLimits(ctx, Limits{Scope: ScopeKey, ScopeID: k.ID, RPM: intp(600)}); err != nil {
		t.Fatal(err)
	}
	s.SetDefaultLimits(intp(60), nil)

	e, err := s.Resolve(ctx, k.ID)
	if err != nil {
		t.Fatal(err)
	}
	if e.RPM == nil || *e.RPM != 600 {
		t.Errorf("RPM = %v; an explicit limit above the default was overridden", e.RPM)
	}
}

func TestNoDefaultsLeavesLimitsUnset(t *testing.T) {
	// The old behaviour, still available deliberately.
	s := openDB(t)
	proj := tree(t, s)
	_, k := mustKey(t, s, "k1", proj)

	e, err := s.Resolve(context.Background(), k.ID)
	if err != nil {
		t.Fatal(err)
	}
	if e.RPM != nil || e.TPM != nil {
		t.Errorf("limits appeared with no defaults set: rpm=%v tpm=%v", e.RPM, e.TPM)
	}
}

func TestControlPlanesAreKeptInTheOrderTheyWereAdded(t *testing.T) {
	ctx := context.Background()
	s := openDB(t)
	at := time.Now()
	for i, u := range []string{"https://10.0.0.140:7443", "https://10.0.0.150:7443"} {
		if err := s.PutControlPlane(ctx, ControlPlane{URL: u, Name: "c", AddedAt: at.Add(time.Duration(i) * time.Second)}); err != nil {
			t.Fatal(err)
		}
	}
	// Recorded again, as init and server add do: not moved, not doubled.
	if err := s.PutControlPlane(ctx, ControlPlane{URL: "https://10.0.0.140:7443", Name: "ctl"}); err != nil {
		t.Fatal(err)
	}
	cps, err := s.ControlPlanes(ctx)
	if err != nil || len(cps) != 2 || cps[0].URL != "https://10.0.0.140:7443" || cps[0].Name != "ctl" {
		t.Fatalf("%+v %v", cps, err)
	}
	if err := s.RemoveControlPlane(ctx, cps[1].URL); err != nil {
		t.Fatal(err)
	}
	if cps, _ := s.ControlPlanes(ctx); len(cps) != 1 {
		t.Errorf("%+v", cps)
	}
}