llmgw/internal/store/limits_owner_test.go

package store

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

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

	for _, err := range []error{
		s.PutOrg(ctx, Org{ID: "o", Name: "o"}),
		s.PutTeam(ctx, Team{ID: "t", OrgID: "o", Name: "t"}),
		s.PutProject(ctx, Project{ID: "p", TeamID: "t", Name: "p"}),
	} {
		if err != nil {
			t.Fatal(err)
		}
	}
	_, k, _ := NewKey("k1", "n", "p")
	k.CreatedAt = time.Now()
	if err := s.CreateKey(ctx, k); err != nil {
		t.Fatal(err)
	}
	return s
}

func TestAResolvedLimitRemembersWhichScopeSetIt(t *testing.T) {
	s := treeStore(t)
	rpm := 60
	if err := s.PutLimits(t.Context(), Limits{Scope: ScopeTeam, ScopeID: "t", RPM: &rpm}); err != nil {
		t.Fatal(err)
	}

	e, err := s.Resolve(t.Context(), "k1")
	if err != nil {
		t.Fatal(err)
	}
	// Without this the tree hands down values and never pools spending: the
	// team's 60 becomes 60 for each key under it.
	if e.RPMOwner.Scope != ScopeTeam || e.RPMOwner.ID != "t" {
		t.Errorf("expected the team to own the rate limit, got %+v", e.RPMOwner)
	}
	if e.RPMOwner.CounterKey() != "team:t" {
		t.Errorf("counter key should name the level too; got %q", e.RPMOwner.CounterKey())
	}
}

func TestTheTighterScopeTakesOwnershipWithTheValue(t *testing.T) {
	s := treeStore(t)
	loose, tight := 100, 10
	must := func(err error) {
		t.Helper()
		if err != nil {
			t.Fatal(err)
		}
	}
	must(s.PutLimits(t.Context(), Limits{Scope: ScopeTeam, ScopeID: "t", RPM: &loose}))
	must(s.PutLimits(t.Context(), Limits{Scope: ScopeProject, ScopeID: "p", RPM: &tight}))

	e, err := s.Resolve(t.Context(), "k1")
	if err != nil {
		t.Fatal(err)
	}
	// The value and the counter have to come from the same place. Enforcing
	// the project's 10 against the team's pool would let one project's traffic
	// exhaust a limit another project never agreed to.
	if *e.RPM != 10 || e.RPMOwner.ID != "p" {
		t.Errorf("expected the project's value and ownership, got %d from %+v", *e.RPM, e.RPMOwner)
	}
}

func TestALooserInnerScopeDoesNotTakeOwnership(t *testing.T) {
	s := treeStore(t)
	tight, loose := 10, 100
	if err := s.PutLimits(t.Context(), Limits{Scope: ScopeTeam, ScopeID: "t", RPM: &tight}); err != nil {
		t.Fatal(err)
	}
	if err := s.PutLimits(t.Context(), Limits{Scope: ScopeProject, ScopeID: "p", RPM: &loose}); err != nil {
		t.Fatal(err)
	}

	e, err := s.Resolve(t.Context(), "k1")
	if err != nil {
		t.Fatal(err)
	}
	// A project cannot raise what its team set, so the team still owns both the
	// number and the counter it is enforced against.
	if *e.RPM != 10 || e.RPMOwner.ID != "t" {
		t.Errorf("expected the team to keep ownership, got %d from %+v", *e.RPM, e.RPMOwner)
	}
}

func TestAClusterDefaultIsOwnedByTheKey(t *testing.T) {
	s := treeStore(t)
	rpm := 30
	s.SetDefaultLimits(&rpm, nil)

	e, err := s.Resolve(t.Context(), "k1")
	if err != nil {
		t.Fatal(err)
	}
	// A default is a fallback for each caller given no limit of its own. Owned
	// by a shared scope it would become one cluster-wide rate limit that every
	// key in every organization competes for — not a default, an outage.
	if e.RPMOwner.Scope != ScopeKey || e.RPMOwner.ID != "k1" {
		t.Errorf("a default should be counted per key, got %+v", e.RPMOwner)
	}
}

func TestSpendComesBackWithThePlaceInTheTree(t *testing.T) {
	s := treeStore(t)
	if err := s.WriteUsage(t.Context(), []Usage{
		{At: time.Now(), KeyID: "k1", ProjectID: "p", Model: "m",
			ServedModel: "m", CostMicros: 700, Status: 200},
		{At: time.Now(), KeyID: "k1", ProjectID: "p", Model: "m",
			ServedModel: "m", CostMicros: 300, Status: 200},
	}); err != nil {
		t.Fatal(err)
	}

	rows, err := s.SpendSinceByScope(t.Context(), time.Now().Add(-time.Hour))
	if err != nil {
		t.Fatal(err)
	}
	if len(rows) != 1 {
		t.Fatalf("expected one key's total, got %+v", rows)
	}
	// The tree, not just the key: a budget owned by a team is rebuilt by
	// summing every key under it, and per-key totals alone cannot say which
	// keys those are.
	got := rows[0]
	if got.Micros != 1000 {
		t.Errorf("expected the two rows summed, got %d", got.Micros)
	}
	if got.ProjectID != "p" || got.TeamID != "t" || got.OrgID != "o" {
		t.Errorf("the whole chain should come back, got %+v", got)
	}
}