llmgw/internal/web/admin_quotas_test.go

package web

import (
	"context"
	"io"
	"log/slog"
	"net/http"
	"net/url"
	"strings"
	"testing"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

type fakeQuotas struct {
	orgs     []store.Org
	teams    []store.TeamPath
	projects []store.ProjectPath
	limits   map[string]store.Limits
	users    []store.User
	usage    map[store.Scope][]store.ScopeUsage

	put     []store.Limits
	cleared [][2]string
}

func (f *fakeQuotas) AllOrgs(context.Context) ([]store.Org, error) { return f.orgs, nil }
func (f *fakeQuotas) AllTeams(context.Context) ([]store.TeamPath, error) {
	return f.teams, nil
}

func (f *fakeQuotas) AllProjects(context.Context) ([]store.ProjectPath, error) {
	return f.projects, nil
}

func (f *fakeQuotas) AllLimits(context.Context) (map[string]store.Limits, error) {
	if f.limits == nil {
		f.limits = map[string]store.Limits{}
	}
	return f.limits, nil
}

func (f *fakeQuotas) PutLimits(_ context.Context, l store.Limits) error {
	f.put = append(f.put, l)
	return nil
}

func (f *fakeQuotas) ClearLimits(_ context.Context, s store.Scope, id string) error {
	f.cleared = append(f.cleared, [2]string{string(s), id})
	return nil
}

func (f *fakeQuotas) UsageByScope(_ context.Context, s store.Scope, _ time.Time) ([]store.ScopeUsage, error) {
	return f.usage[s], nil
}

func quotasPage(t *testing.T, s *fakeQuotas, admin bool) *http.ServeMux {
	t.Helper()
	a := fakeAccounts{user: store.User{ID: "u1", Email: "a@example.com", IsAdmin: admin}}
	p := NewAccountPages(a, nil, nil, "v0", slog.New(slog.NewTextHandler(io.Discard, nil)))
	mux := http.NewServeMux()
	p.RegisterAdminQuotas(mux, s, pricedYes{})
	return mux
}

func i64(v int64) *int64 { return &v }

// pricedYes stands for a cluster where models cost something, so the warning
// about budgets binding on nothing stays out of the way of the other tests.
type pricedYes struct{}

func (pricedYes) AnythingPriced() bool { return true }

type pricedNo struct{}

func (pricedNo) AnythingPriced() bool { return false }

// oneTeam is an org with a team and a project under it.
func (f *fakeQuotas) Users(context.Context, store.UserState) ([]store.User, error) {
	return f.users, nil
}

func (f *fakeQuotas) UsageByUser(context.Context, time.Time) ([]store.UserUsage, error) {
	return nil, nil
}

// Tree derives the tree from the flat lists the fake is given.
func (f *fakeQuotas) Tree(context.Context) ([]store.TreeOrg, error) {
	var out []store.TreeOrg
	for _, o := range f.orgs {
		org := store.TreeOrg{ID: o.ID, Name: o.Name}
		for _, t := range f.teams {
			if t.OrgID != o.ID {
				continue
			}
			team := store.TreeTeam{ID: t.ID, Name: t.Name}
			for _, p := range f.projects {
				if prefix := o.Name + " / " + t.Name + " / "; strings.HasPrefix(p.Path, prefix) {
					team.Projects = append(team.Projects,
						store.TreeProject{ID: p.ID, Name: strings.TrimPrefix(p.Path, prefix)})
				}
			}
			org.Teams = append(org.Teams, team)
		}
		out = append(out, org)
	}
	return out, nil
}

func oneTeam() *fakeQuotas {
	return &fakeQuotas{
		orgs:     []store.Org{{ID: "org-1", Name: "Scania"}},
		teams:    []store.TeamPath{{ID: "team-1", Name: "platform", OrgID: "org-1", OrgName: "Scania"}},
		projects: []store.ProjectPath{{ID: "proj-1", Path: "Scania / platform / gateway"}},
		limits:   map[string]store.Limits{},
	}
}

func TestTheWholeTreeIsListedWhetherOrNotItHasAQuota(t *testing.T) {
	body := mirrorGet(t, quotasPage(t, oneTeam(), true), "/admin/quotas").Body.String()

	// A group with no quota is the one most likely to need one, so it cannot
	// be the one the page leaves out.
	for _, want := range []string{"Scania", "platform", "gateway"} {
		if !strings.Contains(body, want) {
			t.Errorf("%q is missing from the page", want)
		}
	}
}

func TestSpendIsShownAgainstTheBudget(t *testing.T) {
	s := oneTeam()
	s.limits["team:team-1"] = store.Limits{
		Scope: store.ScopeTeam, ScopeID: "team-1",
		Budget: i64(2_000_000), BudgetPeriod: store.PeriodMonth,
	}
	s.usage = map[store.Scope][]store.ScopeUsage{
		store.ScopeTeam: {{ID: "team-1", CostMicros: 1_000_000}},
	}

	body := mirrorGet(t, quotasPage(t, s, true), "/admin/quotas").Body.String()
	if !strings.Contains(body, "50%") {
		t.Errorf("expected half the budget used to be shown; got %s", body)
	}
}

func TestALimitAnAncestorAlreadyCapsIsMarked(t *testing.T) {
	s := oneTeam()
	// The team is given more than its organization allows. Resolution takes the
	// most restrictive, so this is not a bigger budget — it is a number that
	// will never be reached.
	s.limits["org:org-1"] = store.Limits{
		Scope: store.ScopeOrg, ScopeID: "org-1", Budget: i64(1_000_000),
		BudgetPeriod: store.PeriodMonth,
	}
	s.limits["team:team-1"] = store.Limits{
		Scope: store.ScopeTeam, ScopeID: "team-1", Budget: i64(5_000_000),
		BudgetPeriod: store.PeriodMonth,
	}

	body := mirrorGet(t, quotasPage(t, s, true), "/admin/quotas").Body.String()
	if !strings.Contains(body, "Scania caps this at") {
		t.Errorf("a shadowed limit should say so; got %s", body)
	}
}

func TestALimitNothingCapsIsNotMarked(t *testing.T) {
	s := oneTeam()
	s.limits["org:org-1"] = store.Limits{
		Scope: store.ScopeOrg, ScopeID: "org-1", Budget: i64(9_000_000),
		BudgetPeriod: store.PeriodMonth,
	}
	s.limits["team:team-1"] = store.Limits{
		Scope: store.ScopeTeam, ScopeID: "team-1", Budget: i64(1_000_000),
		BudgetPeriod: store.PeriodMonth,
	}

	body := mirrorGet(t, quotasPage(t, s, true), "/admin/quotas").Body.String()
	if strings.Contains(body, "caps this at") {
		t.Error("a limit stricter than its parent's is in force and must not be flagged")
	}
}

func TestABudgetIsTakenInWholeUnitsAndStoredInMicros(t *testing.T) {
	s := oneTeam()
	mirrorPost(t, quotasPage(t, s, true), "/admin/quotas", url.Values{
		"csrf": {"tok"}, "scope": {"team"}, "id": {"team-1"},
		"budget": {"1000"}, "period": {"month"},
	})

	if len(s.put) != 1 || s.put[0].Budget == nil || *s.put[0].Budget != 1_000_000_000 {
		t.Fatalf("expected 1000 units stored as micro-units, got %+v", s.put)
	}
}

func TestABudgetWithoutAPeriodIsRefused(t *testing.T) {
	s := oneTeam()
	body := mirrorPost(t, quotasPage(t, s, true), "/admin/quotas", url.Values{
		"csrf": {"tok"}, "scope": {"team"}, "id": {"team-1"}, "budget": {"1000"},
	}).Body.String()

	// A monthly budget credited to a daily counter is a budget that never
	// binds, so the period cannot be left to a default.
	if len(s.put) != 0 {
		t.Error("nothing should have been written")
	}
	if !strings.Contains(body, "period") {
		t.Errorf("the error should name what is missing; got %s", body)
	}
}

func TestABlankFieldMeansNoLimitRatherThanZero(t *testing.T) {
	s := oneTeam()
	mirrorPost(t, quotasPage(t, s, true), "/admin/quotas", url.Values{
		"csrf": {"tok"}, "scope": {"team"}, "id": {"team-1"}, "rpm": {"60"},
	})

	// Zero would be a limit of nothing — every request refused — which is the
	// opposite of what an empty box means.
	if len(s.put) != 1 {
		t.Fatalf("expected one write, got %+v", s.put)
	}
	if s.put[0].TPM != nil {
		t.Errorf("a blank field should leave no limit, got %d", *s.put[0].TPM)
	}
	if s.put[0].RPM == nil || *s.put[0].RPM != 60 {
		t.Errorf("the value given should be written, got %+v", s.put[0].RPM)
	}
}

func TestRubbishInANumberIsRefused(t *testing.T) {
	s := oneTeam()
	mirrorPost(t, quotasPage(t, s, true), "/admin/quotas", url.Values{
		"csrf": {"tok"}, "scope": {"team"}, "id": {"team-1"}, "rpm": {"lots"},
	})
	if len(s.put) != 0 {
		t.Error("nothing should have been written")
	}
}

func TestAQuotaCanBeRemovedEntirely(t *testing.T) {
	s := oneTeam()
	mirrorPost(t, quotasPage(t, s, true), "/admin/quotas", url.Values{
		"csrf": {"tok"}, "scope": {"team"}, "id": {"team-1"}, "clear": {"1"},
	})

	// A row of nulls and no row resolve the same, but only one of them tells
	// somebody reading the table that nobody has set a limit here.
	if len(s.cleared) != 1 || s.cleared[0] != [2]string{"team", "team-1"} {
		t.Errorf("expected the limits to be cleared, got %v", s.cleared)
	}
	if len(s.put) != 0 {
		t.Error("clearing must not also write a row")
	}
}

func TestAQuotaOnSomethingThatIsNotAGroupIsRefused(t *testing.T) {
	s := oneTeam()
	body := mirrorPost(t, quotasPage(t, s, true), "/admin/quotas", url.Values{
		"csrf": {"tok"}, "scope": {"key"}, "id": {"key-1"}, "rpm": {"1"},
	}).Body.String()

	// Per-key limits stay on the CLI: an admin session reaching individual keys
	// would put every caller's throughput one click away.
	if len(s.put) != 0 {
		t.Error("a key limit should not be settable here")
	}
	if !strings.Contains(body, "organization") {
		t.Errorf("the refusal should say what is allowed; got %s", body)
	}
}

func TestChangingAQuotaNeedsTheCSRFToken(t *testing.T) {
	s := oneTeam()
	w := mirrorPost(t, quotasPage(t, s, true), "/admin/quotas", url.Values{
		"scope": {"team"}, "id": {"team-1"}, "rpm": {"60"},
	})
	if w.Code != http.StatusBadRequest || len(s.put) != 0 {
		t.Errorf("expected 400 and no write, got %d %+v", w.Code, s.put)
	}
}

func TestSomebodyWhoIsNotAnAdminSeesNoQuotas(t *testing.T) {
	if w := mirrorGet(t, quotasPage(t, oneTeam(), false), "/admin/quotas"); w.Code != http.StatusNotFound {
		t.Errorf("expected 404 for a non-admin, got %d", w.Code)
	}
}

func TestAFormOpensOnTheBranchThatWasClicked(t *testing.T) {
	body := mirrorGet(t, quotasPage(t, oneTeam(), true),
		"/admin/quotas?edit=team:team-1").Body.String()

	// The tree is the interface: a quota belongs to a branch, so the form opens
	// on that branch rather than asking for its id again. Ids are random — two
	// organizations can both have a "platform" team — so typing one is a typo
	// waiting to write a quota onto nothing and report success.
	if !strings.Contains(body, `name="id" value="team-1"`) {
		t.Errorf("the open form should carry the branch it belongs to; got %s", body)
	}
	if !strings.Contains(body, "Scania / platform") {
		t.Error("the form should name the branch in full")
	}
}

func TestOnlyTheClickedBranchOpens(t *testing.T) {
	body := mirrorGet(t, quotasPage(t, oneTeam(), true),
		"/admin/quotas?edit=team:team-1").Body.String()

	// One form at a time. A disclosure on every row would put a form in the
	// markup for each branch of a page whose job is to be read as a tree.
	// The menu's sign-out form carries a token too; it is not a quota form.
	if n := strings.Count(body, `name="csrf"`) - strings.Count(body, `action="/logout"`); n != 1 {
		t.Errorf("expected exactly one form, found %d", n)
	}
	if !strings.Contains(body, `href="/admin/quotas?edit=project:proj-1"`) {
		t.Error("every other branch should still offer to open")
	}
}

func TestABranchWithNoQuotaSaysSoOnItsLink(t *testing.T) {
	body := mirrorGet(t, quotasPage(t, oneTeam(), true), "/admin/quotas").Body.String()

	// "change" on a branch that has nothing set would suggest there is
	// something to change.
	if !strings.Contains(body, `">set</a>`) {
		t.Errorf("an unlimited branch should invite setting one; got %s", body)
	}
}

func TestSavingClosesTheFormSoTheResultIsVisible(t *testing.T) {
	s := oneTeam()
	body := mirrorPost(t, quotasPage(t, s, true), "/admin/quotas", url.Values{
		"csrf": {"tok"}, "scope": {"team"}, "id": {"team-1"}, "rpm": {"60"},
	}).Body.String()

	// Re-rendering the form you just submitted hides the row you changed behind
	// it, and invites a second submission of the same thing.
	if strings.Contains(body, `name="rpm"`) {
		t.Error("the form should be closed after a save")
	}
	if !strings.Contains(body, "rpm 60") {
		t.Errorf("the page should confirm what was written; got %s", body)
	}
}

func TestABudgetThatCannotBindSaysSo(t *testing.T) {
	a := fakeAccounts{user: store.User{ID: "u1", Email: "a@example.com", IsAdmin: true}}
	p := NewAccountPages(a, nil, nil, "v0", slog.New(slog.NewTextHandler(io.Discard, nil)))
	mux := http.NewServeMux()
	p.RegisterAdminQuotas(mux, oneTeam(), pricedNo{})

	// Whitespace normalised: the sentence wraps in the template and a test
	// that pinned the wrap would break on a reflow rather than on a change of
	// meaning.
	body := strings.Join(strings.Fields(mirrorGet(t, mux, "/admin/quotas").Body.String()), " ")

	// A budget is spend per period and spend comes from a price table. With no
	// prices the number reads as a limit, reports zero percent used, and binds
	// on nothing.
	if !strings.Contains(body, "budgets stay at zero") {
		t.Errorf("the page should say budgets cannot move; got %s", body)
	}
	// Said precisely: the other two limits are unaffected and still work.
	if !strings.Contains(body, "Rate and token limits still apply") {
		t.Error("the warning should not read as though no quota works")
	}
}

func TestAPricedClusterIsNotWarned(t *testing.T) {
	body := mirrorGet(t, quotasPage(t, oneTeam(), true), "/admin/quotas").Body.String()
	if strings.Contains(body, "stays at zero") {
		t.Error("a cluster with prices should see no warning")
	}
}

func TestABranchShowsTheLimitItIsHeldToFromAbove(t *testing.T) {
	q := oneTeam()
	q.limits[store.LimitOwner{Scope: store.ScopeOrg, ID: "org-1"}.CounterKey()] =
		store.Limits{Scope: store.ScopeOrg, ScopeID: "org-1", RPM: ptrInt(60)}
	body := mirrorGet(t, quotasPage(t, q, true), "/admin/quotas").Body.String()

	// The project sets nothing, and read "—" — which looked unlimited while
	// its organization held it to sixty a minute.
	if n := strings.Count(body, `<span class="q-up" title="set on Scania">60 ↑</span>`); n != 2 {
		t.Errorf("the team and the project should both show the organization's 60; found %d", n)
	}
}

func TestALargerLimitBelowSaysItIsCapped(t *testing.T) {
	q := oneTeam()
	q.limits[store.LimitOwner{Scope: store.ScopeOrg, ID: "org-1"}.CounterKey()] =
		store.Limits{Scope: store.ScopeOrg, ScopeID: "org-1", RPM: ptrInt(60)}
	q.limits[store.LimitOwner{Scope: store.ScopeProject, ID: "proj-1"}.CounterKey()] =
		store.Limits{Scope: store.ScopeProject, ScopeID: "proj-1", RPM: ptrInt(500)}
	body := mirrorGet(t, quotasPage(t, q, true), "/admin/quotas").Body.String()
	if !strings.Contains(body, "Scania caps this at 60 requests/min") {
		t.Error("a limit that can never be reached should say what holds it")
	}
}

func ptrInt(v int) *int { return &v }

func TestAPerMemberAndAPersonsLimitAreSavedWhereTheyBelong(t *testing.T) {
	q := oneTeam()
	q.users = []store.User{{ID: "usr-a", Email: "ayse@x", State: store.UserActive}}
	mux := quotasPage(t, q, true)

	mirrorPost(t, mux, "/admin/quotas", url.Values{"csrf": {"tok"},
		"scope": {"project-member"}, "id": {"proj-1"}, "rpm": {"20"}})
	mirrorPost(t, mux, "/admin/quotas", url.Values{"csrf": {"tok"},
		"scope": {"user"}, "id": {"usr-a"}, "budget": {"50"}, "period": {"month"}})

	if len(q.put) != 2 {
		t.Fatalf("saved %d limits, want 2", len(q.put))
	}
	if q.put[0].Scope != store.ScopeProjectMember || q.put[0].ScopeID != "proj-1" || *q.put[0].RPM != 20 {
		t.Errorf("per-member limit saved as %+v", q.put[0])
	}
	if q.put[1].Scope != store.ScopeUser || q.put[1].ScopeID != "usr-a" || *q.put[1].Budget != 50_000_000 {
		t.Errorf("person's limit saved as %+v", q.put[1])
	}
}

func TestThePeopleSectionShowsEachPersonsOwnLimit(t *testing.T) {
	q := oneTeam()
	q.users = []store.User{{ID: "usr-a", Email: "ayse@x", State: store.UserActive}}
	q.limits[store.LimitOwner{Scope: store.ScopeUser, ID: "usr-a"}.CounterKey()] =
		store.Limits{Scope: store.ScopeUser, ScopeID: "usr-a", RPM: ptrInt(7)}
	q.limits[store.LimitOwner{Scope: store.ScopeProjectMember, ID: "proj-1"}.CounterKey()] =
		store.Limits{Scope: store.ScopeProjectMember, ScopeID: "proj-1", RPM: ptrInt(20)}
	body := strings.Join(strings.Fields(
		mirrorGet(t, quotasPage(t, q, true), "/admin/quotas").Body.String()), " ")
	if !strings.Contains(body, `ayse@x</span>`) || !strings.Contains(body, `<td class="num">7</td>`) {
		t.Error("the person and their limit are not on the page")
	}
	if !strings.Contains(body, "each member here: rpm 20") {
		t.Error("the project's per-member limit is not shown")
	}
}