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
	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)
	return mux
}

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

// oneTeam is an org with a team and a project under it.
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.
	if n := strings.Count(body, `name="csrf"`); 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 quota") {
		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)
	}
}