llmgw/internal/store/postgres_test.go

package store

import (
	"context"
	"errors"
	"net/url"
	"os"
	"strings"
	"testing"
	"time"
)

// openPG opens the test database named by LLMGW_TEST_POSTGRES, or skips.
//
// Env-gated rather than spun up here, because a store test that needs a
// container is a test nobody runs. Pointed at a database, it runs the same
// assertions the SQLite tests make — which is the only thing that can catch
// the two backends drifting apart.
//
// Every table is emptied first. The alternative is a test database that
// accumulates rows until an assertion about counts starts failing for reasons
// that have nothing to do with the change being tested.
func openPG(t *testing.T) *DB {
	t.Helper()
	dsn := os.Getenv("LLMGW_TEST_POSTGRES")
	if dsn == "" {
		t.Skip("set LLMGW_TEST_POSTGRES to a database DSN to run these")
	}
	requireTestDatabase(t, dsn)

	s, err := OpenPostgres(context.Background(), dsn)
	if err != nil {
		t.Fatalf("opening the test database: %v", err)
	}
	t.Cleanup(func() { _ = s.Close() })

	// CASCADE because the tables reference each other, and TRUNCATE rather
	// than DELETE so identity columns restart too.
	_, err = s.db.ExecContext(context.Background(),
		// leadership included, and it is not in copyOrder: a subtest that
		// inherits the previous one's lease is refused for reasons that have
		// nothing to do with what it is testing. It passed on SQLite, where
		// every subtest gets a fresh file, and failed only here.
		`TRUNCATE thread_messages, threads, memory_notes, usage, api_keys, memberships,
		         key_claims, limits, models, sessions, users, projects, teams,
		         organizations, leadership RESTART IDENTITY CASCADE`)
	if err != nil {
		t.Fatalf("clearing the test database: %v", err)
	}
	return s
}

// requireTestDatabase refuses to run against a database that is not obviously
// a test one.
//
// These tests truncate every table before each one. Pointed at a real
// deployment that is a silent, complete data loss — every key, every user,
// every conversation — and it is not a hypothetical: it happened here, through
// an SSH tunnel to the production database that had been opened minutes
// earlier for a perfectly good reason. The DSN was correct for what it was
// opened for and catastrophic for what it was then used for.
//
// The name has to say so. A guard that could be satisfied by accident would
// not have stopped it.
func requireTestDatabase(t *testing.T, dsn string) {
	t.Helper()
	u, err := url.Parse(dsn)
	if err != nil {
		t.Fatalf("LLMGW_TEST_POSTGRES is not a URL: %v", err)
	}
	name := strings.TrimPrefix(u.Path, "/")
	if !strings.Contains(strings.ToLower(name), "test") {
		t.Fatalf("refusing to run destructive tests against database %q:\n"+
			"  these TRUNCATE every table before each test.\n"+
			"  Point LLMGW_TEST_POSTGRES at a database whose name contains \"test\".", name)
	}
}

func TestPostgresCarriesTheWholeTenancyTree(t *testing.T) {
	ctx := context.Background()
	s := openPG(t)

	proj := tree(t, s)
	secret, k := mustKey(t, s, "k1", proj)
	_ = secret

	got, err := s.KeyByHash(ctx, k.Hash)
	if err != nil {
		t.Fatalf("reading the key back: %v", err)
	}
	if got.ID != k.ID || got.ProjectID != proj {
		t.Errorf("key came back as %+v", got)
	}
}

func TestPostgresKeepsUsageAndSpend(t *testing.T) {
	// The accounting path, which is the one that decides whether a budget
	// means anything after a restart.
	ctx := context.Background()
	s := openPG(t)
	proj := tree(t, s)
	_, k := mustKey(t, s, "k1", proj)

	now := time.Now()
	err := s.WriteUsage(ctx, []Usage{
		{At: now, KeyID: k.ID, ProjectID: proj, Model: "m", ServedModel: "m",
			PromptTokens: 10, CompletionTokens: 20, CostMicros: 500, Status: 200, LatencyMs: 42},
		{At: now, KeyID: k.ID, ProjectID: proj, Model: "m", ServedModel: "m",
			PromptTokens: 1, CompletionTokens: 2, CostMicros: 250, Status: 200, LatencyMs: 7},
	})
	if err != nil {
		t.Fatalf("writing usage: %v", err)
	}

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

func TestPostgresKeepsThreadsInOrderAndScoped(t *testing.T) {
	// The newest tables, and the ones with the most SQL per operation —
	// a transaction, a MAX(seq), an insert per turn.
	ctx := context.Background()
	s := openPG(t)
	proj := tree(t, s)
	if err := s.SetThreadRetention(ctx, proj, 24); err != nil {
		t.Fatal(err)
	}

	for _, m := range []Message{
		{Role: "user", Content: "first"},
		{Role: "assistant", Content: "answer", Model: "small"},
		{Role: "user", Content: "second"},
	} {
		if err := s.AppendThread(ctx, proj, "th1", []Message{m}); err != nil {
			t.Fatal(err)
		}
	}

	got, err := s.Thread(ctx, proj, "th1")
	if err != nil {
		t.Fatal(err)
	}
	if len(got) != 3 || got[0].Content != "first" || got[2].Content != "second" {
		t.Fatalf("thread came back as %+v", got)
	}
	if got[1].Model != "small" {
		t.Errorf("the model that answered was lost: %+v", got[1])
	}

	// And the scope holds here exactly as it does on SQLite.
	if err := s.PutProject(ctx, Project{ID: "proj2", TeamID: "team1", Name: "Other"}); err != nil {
		t.Fatal(err)
	}
	if _, err := s.Thread(ctx, "proj2", "th1"); !errors.Is(err, ErrNoThread) {
		t.Errorf("reading across projects returned %v", err)
	}
}

func TestPostgresEnforcesLimitsTheSameWay(t *testing.T) {
	ctx := context.Background()
	s := openPG(t)
	proj := tree(t, s)
	_, k := mustKey(t, s, "k1", proj)

	rpm := 7
	budget := int64(9000)
	if err := s.PutLimits(ctx, Limits{
		Scope: ScopeProject, ScopeID: proj, RPM: &rpm,
		Budget: &budget, BudgetPeriod: PeriodMonth,
	}); err != nil {
		t.Fatal(err)
	}

	eff, err := s.Resolve(ctx, k.ID)
	if err != nil {
		t.Fatal(err)
	}
	if eff.RPM == nil || *eff.RPM != 7 {
		t.Errorf("rpm did not resolve down the tree: %+v", eff.RPM)
	}
	if eff.Budget == nil || *eff.Budget != 9000 {
		t.Errorf("budget did not resolve down the tree: %+v", eff.Budget)
	}
}

func TestPostgresUpsertsRatherThanDuplicating(t *testing.T) {
	// ON CONFLICT is the one non-trivial construct shared by both dialects.
	// If it behaved differently here, applying a configuration twice would
	// start failing instead of being a no-op.
	ctx := context.Background()
	s := openPG(t)
	tree(t, s)
	tree(t, s)

	orgs, err := s.Orgs(ctx)
	if err != nil {
		t.Fatal(err)
	}
	if len(orgs) != 1 {
		t.Errorf("applying the tree twice produced %d orgs", len(orgs))
	}
}

func TestPostgresBucketsUsageByEveryGrouping(t *testing.T) {
	// The query that most deserves this. It selects a computed bucket
	// expression and groups by ordinal, and three earlier reports shipped
	// working on SQLite and failing here because PostgreSQL insists every
	// non-aggregated column is grouped and matches grouped expressions to
	// selected ones by their text.
	ctx := context.Background()
	s := openPG(t)
	proj := tree(t, s)
	_, k := mustKey(t, s, "k1", proj)

	now := time.Now()
	err := s.WriteUsage(ctx, []Usage{
		{At: now.Add(-2 * time.Hour), KeyID: k.ID, ProjectID: proj,
			Model: "m", ServedModel: "m",
			PromptTokens: 100, CompletionTokens: 50, CostMicros: 900, Status: 200},
		{At: now, KeyID: k.ID, ProjectID: proj, Model: "m", ServedModel: "m",
			PromptTokens: 10, CompletionTokens: 5, CostMicros: 100, Status: 500},
	})
	if err != nil {
		t.Fatalf("writing usage: %v", err)
	}

	for _, g := range Groupings() {
		rows, err := s.UsageBucketsBy(ctx, g, now.Add(-6*time.Hour), time.Hour)
		if err != nil {
			t.Fatalf("grouping by %s: %v", g, err)
		}
		if len(rows) != 1 {
			t.Fatalf("grouping by %s returned %d groups, want 1", g, len(rows))
		}
		if rows[0].CostMicros != 1000 || rows[0].Requests != 2 || rows[0].Errors != 1 {
			t.Errorf("grouping by %s: %d micros, %d requests, %d errors",
				g, rows[0].CostMicros, rows[0].Requests, rows[0].Errors)
		}
		busy := 0
		for _, b := range rows[0].Buckets {
			if b.Requests > 0 {
				busy++
			}
		}
		if busy != 2 {
			t.Errorf("grouping by %s put the traffic in %d buckets, want 2", g, busy)
		}
	}
}