llmgw/internal/store/usagescope_test.go

package store

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

// twoTeams is one org, two teams, a project each, a key each.
func twoTeams(t *testing.T) *DB {
	t.Helper()
	ctx := context.Background()
	s, err := OpenSQLite(ctx, filepath.Join(t.TempDir(), "u.db"))
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = s.Close() })

	must := func(err error) {
		t.Helper()
		if err != nil {
			t.Fatal(err)
		}
	}
	must(s.PutOrg(ctx, Org{ID: "o", Name: "Acme"}))
	must(s.PutTeam(ctx, Team{ID: "t1", OrgID: "o", Name: "platform"}))
	must(s.PutTeam(ctx, Team{ID: "t2", OrgID: "o", Name: "research"}))
	must(s.PutProject(ctx, Project{ID: "p1", TeamID: "t1", Name: "gateway"}))
	must(s.PutProject(ctx, Project{ID: "p2", TeamID: "t2", Name: "labs"}))
	for id, proj := range map[string]string{"k1": "p1", "k2": "p1", "k3": "p2"} {
		_, k, _ := NewKey(id, "n", proj)
		k.CreatedAt = time.Now()
		must(s.CreateKey(ctx, k))
	}
	must(s.WriteUsage(ctx, []Usage{
		{At: time.Now(), KeyID: "k1", ProjectID: "p1", Model: "m", ServedModel: "m",
			PromptTokens: 10, CompletionTokens: 5, CostMicros: 600, Status: 200},
		{At: time.Now(), KeyID: "k2", ProjectID: "p1", Model: "m", ServedModel: "m",
			PromptTokens: 10, CompletionTokens: 5, CostMicros: 400, Status: 500},
		{At: time.Now(), KeyID: "k3", ProjectID: "p2", Model: "m", ServedModel: "m",
			PromptTokens: 1, CompletionTokens: 1, CostMicros: 100, Status: 200},
	}))
	return s
}

func byID(rows []ScopeUsage, id string) (ScopeUsage, bool) {
	for _, r := range rows {
		if r.ID == id {
			return r, true
		}
	}
	return ScopeUsage{}, false
}

func TestATeamsSpendIsTheSumOfItsKeys(t *testing.T) {
	s := twoTeams(t)

	rows, err := s.UsageByScope(t.Context(), ScopeTeam, time.Now().Add(-time.Hour))
	if err != nil {
		t.Fatal(err)
	}
	// Two members of one team, spending separately, is one number for the team.
	// Per-key totals could not answer this at all.
	got, ok := byID(rows, "t1")
	if !ok {
		t.Fatalf("the team is missing from %+v", rows)
	}
	if got.CostMicros != 1000 {
		t.Errorf("expected both members summed, got %d", got.CostMicros)
	}
	if got.Keys != 2 {
		t.Errorf("expected two keys counted, got %d", got.Keys)
	}
	if got.Errors != 1 {
		t.Errorf("expected the one failure counted, got %d", got.Errors)
	}
	if got.Name != "platform" {
		t.Errorf("the group should be named, got %q", got.Name)
	}
}

func TestTeamsAreNotAddedTogether(t *testing.T) {
	s := twoTeams(t)
	rows, _ := s.UsageByScope(t.Context(), ScopeTeam, time.Now().Add(-time.Hour))

	if len(rows) != 2 {
		t.Fatalf("expected a row per team, got %+v", rows)
	}
	other, _ := byID(rows, "t2")
	if other.CostMicros != 100 {
		t.Errorf("the second team's spend leaked, got %d", other.CostMicros)
	}
}

func TestTheOrgRollsUpEveryTeamBeneathIt(t *testing.T) {
	s := twoTeams(t)
	rows, err := s.UsageByScope(t.Context(), ScopeOrg, time.Now().Add(-time.Hour))
	if err != nil {
		t.Fatal(err)
	}
	if len(rows) != 1 || rows[0].CostMicros != 1100 {
		t.Errorf("expected one org totalling 1100, got %+v", rows)
	}
}

func TestTheQuotaComesBackBesideTheSpend(t *testing.T) {
	s := twoTeams(t)
	budget := int64(2000)
	if err := s.PutLimits(t.Context(), Limits{
		Scope: ScopeTeam, ScopeID: "t1", Budget: &budget, BudgetPeriod: PeriodMonth,
	}); err != nil {
		t.Fatal(err)
	}

	rows, _ := s.UsageByScope(t.Context(), ScopeTeam, time.Now().Add(-time.Hour))
	got, _ := byID(rows, "t1")
	// Spend without the quota beside it is a number nobody can act on. Half of
	// a budget is the answer; 1000 micro-units on its own is trivia.
	if got.Budget == nil || *got.Budget != 2000 {
		t.Fatalf("the budget set on the team should come back, got %+v", got.Budget)
	}
	frac, ok := got.Spent()
	if !ok || frac != 0.5 {
		t.Errorf("expected half the budget used, got %v %v", frac, ok)
	}
}

func TestAGroupWithNoBudgetOfItsOwnShowsNone(t *testing.T) {
	s := twoTeams(t)
	budget := int64(9000)
	if err := s.PutLimits(t.Context(), Limits{
		Scope: ScopeOrg, ScopeID: "o", Budget: &budget, BudgetPeriod: PeriodMonth,
	}); err != nil {
		t.Fatal(err)
	}

	rows, _ := s.UsageByScope(t.Context(), ScopeTeam, time.Now().Add(-time.Hour))
	got, _ := byID(rows, "t1")
	// The org's budget bounds this team but is not the team's to spend.
	// Printing 9000 beside each of two teams would read as 18000 available,
	// which is the exact misreading this work exists to remove.
	if got.Budget != nil {
		t.Errorf("a team with no budget of its own should show none, got %d", *got.Budget)
	}
	if _, ok := got.Spent(); ok {
		t.Error("there is no fraction to report without a budget at this level")
	}
}

func TestGroupingBySomethingThatIsNotALevelIsRefused(t *testing.T) {
	s := twoTeams(t)
	// Never interpolated: the grouping picks from a fixed table of levels, so
	// an unknown one is an error rather than a fragment of SQL.
	if _, err := s.UsageByScope(t.Context(), Scope("key; DROP TABLE usage"), time.Now()); err == nil {
		t.Error("expected an unknown grouping to be refused")
	}
}