llmgw/internal/web/admin_usage_test.go

package web

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

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

type fakeScopeSeries struct {
	byGroup map[store.Grouping][]store.UsageSeries
	// stats, when set, is what StatsBy answers for every grouping.
	stats map[string]store.GroupStats
	asked []store.Grouping
	width time.Duration
	since time.Time
}

func (f *fakeScopeSeries) UsageBucketsBy(_ context.Context, g store.Grouping,
	since time.Time, width time.Duration,
) ([]store.UsageSeries, error) {
	f.asked = append(f.asked, g)
	f.width, f.since = width, since
	return f.byGroup[g], nil
}

func (f *fakeScopeSeries) StatsBy(context.Context, store.Grouping, time.Time) (map[string]store.GroupStats, error) {
	return f.stats, nil
}

func usagePage(t *testing.T, s *fakeScopeSeries, priced Priced) *http.ServeMux {
	t.Helper()
	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.RegisterAdminUsage(mux, s, priced)
	return mux
}

// hourly builds a group whose traffic rises over six slices, so a chart drawn
// from it has a shape rather than a dot.
func hourly(g store.Grouping, id, name string, scale int64) store.UsageSeries {
	u := store.UsageSeries{Grouping: g, ID: id, Name: name}
	base := time.Now().Add(-6 * time.Hour).Truncate(time.Hour)
	for i := int64(1); i <= 6; i++ {
		b := store.Bucket{
			Start:      base.Add(time.Duration(i) * time.Hour),
			Requests:   i * scale,
			Prompt:     i * scale * 100,
			Completion: i * scale * 20,
			CostMicros: i * scale * 1000,
		}
		u.Buckets = append(u.Buckets, b)
		u.Requests += b.Requests
		u.Prompt += b.Prompt
		u.Completion += b.Completion
		u.CostMicros += b.CostMicros
	}
	return u
}

func TestTheUsagePageDrawsAndTabulatesEachGroup(t *testing.T) {
	s := &fakeScopeSeries{byGroup: map[store.Grouping][]store.UsageSeries{
		store.GroupByTeam: {
			hourly(store.GroupByTeam, "t1", "platform", 3),
			hourly(store.GroupByTeam, "t2", "research", 1),
		},
	}}
	body := mirrorGet(t, usagePage(t, s, pricedYes{}), "/admin/usage").Body.String()

	for _, want := range []string{"platform", "research"} {
		if !strings.Contains(body, want) {
			t.Errorf("%q is missing from the page", want)
		}
	}
	// A chart is an SVG with lines in it, not a frame saying there is nothing.
	if strings.Count(body, "<svg") < 3 {
		t.Errorf("expected a chart for requests, tokens and spend; got %s", body)
	}
	if strings.Contains(body, "Nothing recorded in this window") {
		t.Error("the table was left out although there are two groups")
	}
}

func TestTeamIsTheDefaultGrouping(t *testing.T) {
	// The middle of the tree: an organization is usually one on a first
	// cluster, and a project list is long before it is informative.
	s := &fakeScopeSeries{}
	mirrorGet(t, usagePage(t, s, pricedYes{}), "/admin/usage")

	if len(s.asked) != 1 || s.asked[0] != store.GroupByTeam {
		t.Errorf("asked for %v, expected one query grouped by team", s.asked)
	}
}

func TestTheGroupingComesFromTheURL(t *testing.T) {
	s := &fakeScopeSeries{}
	mirrorGet(t, usagePage(t, s, pricedYes{}), "/admin/usage?by=user")

	if len(s.asked) != 1 || s.asked[0] != store.GroupByUser {
		t.Errorf("asked for %v, expected a query grouped by user", s.asked)
	}
}

func TestAnUnknownGroupingFallsBackRatherThanFailing(t *testing.T) {
	// The grouping reaches SQL, so it is validated before it is used. A page
	// that 500s on a mistyped query string is worse than one that shows the
	// default, and the chooser makes the mistake obvious.
	s := &fakeScopeSeries{}
	rec := mirrorGet(t, usagePage(t, s, pricedYes{}), "/admin/usage?by=drop+table")

	if rec.Code != 200 {
		t.Fatalf("status %d", rec.Code)
	}
	if len(s.asked) != 1 || s.asked[0] != store.GroupByTeam {
		t.Errorf("asked for %v, expected the default", s.asked)
	}
}

func TestTheWindowChoosesItsOwnBucketWidth(t *testing.T) {
	// A month in twenty-minute slices is two thousand points across seven
	// hundred pixels, which is a smear rather than a chart.
	s := &fakeScopeSeries{}
	mirrorGet(t, usagePage(t, s, pricedYes{}), "/admin/usage?window=30d")
	wide := s.width

	s2 := &fakeScopeSeries{}
	mirrorGet(t, usagePage(t, s2, pricedYes{}), "/admin/usage?window=24h")

	if !(wide > s2.width) {
		t.Errorf("a 30-day window bucketed at %s, a 24-hour one at %s", wide, s2.width)
	}
	if s2.width > time.Hour {
		t.Errorf("a day in %s slices is too coarse to read", s2.width)
	}
}

func TestWithNoPricesTheSpendChartIsReplacedBySayingWhy(t *testing.T) {
	s := &fakeScopeSeries{byGroup: map[store.Grouping][]store.UsageSeries{
		store.GroupByTeam: {hourly(store.GroupByTeam, "t1", "platform", 3)},
	}}
	body := mirrorGet(t, usagePage(t, s, pricedNo{}), "/admin/usage").Body.String()

	// A spend chart flat along the bottom reads as a quiet month rather than as
	// a cluster where nothing is priced.
	if !strings.Contains(body, "No model has a price") {
		t.Errorf("the page should say why spend is empty; got %s", body)
	}
	if strings.Count(body, "<svg") > 2 {
		t.Error("a spend chart was drawn although nothing is priced")
	}
}

func TestAPersonWithNoAccountIsNamedRatherThanBlank(t *testing.T) {
	// Keys issued by an operator before accounts existed carry real traffic,
	// and an unlabelled line in a legend is the one nobody can act on.
	s := &fakeScopeSeries{byGroup: map[store.Grouping][]store.UsageSeries{
		store.GroupByUser: {hourly(store.GroupByUser, "", "", 2)},
	}}
	body := mirrorGet(t, usagePage(t, s, pricedYes{}), "/admin/usage?by=user").Body.String()

	if !strings.Contains(body, "keys with no owner") {
		t.Errorf("unowned traffic was left unlabelled; got %s", body)
	}
}

func TestTheUsagePageCountsTheGPUTimeEachOwnerHeld(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)))
	// Over a day in 20-minute steps: team:research held 8 cards throughout,
	// nobody's model 2 for the last three steps.
	res, none := promq.Series{Name: "team:research"}, promq.Series{Name: ""}
	for i := range 72 {
		res.Points = append(res.Points, promq.Point{Value: 8})
		if i >= 69 {
			none.Points = append(none.Points, promq.Point{Value: 2})
		}
	}
	h := &fakeHistory{byExpr: map[string][]promq.Series{`sum by (owner) (llmgw_gpu_assigned)`: {res, none}}}
	p.SetHistory(h)
	mux := http.NewServeMux()
	p.RegisterAdminUsage(mux, &fakeScopeSeries{}, nil)
	body := mirrorGet(t, mux, "/admin/usage?window=24h").Body.String()
	for _, want := range []string{"GPU time, by owner", "<td>team:research</td><td class=\"num\">192.0</td>",
		"<td>no owner</td><td class=\"num\">2.0</td>"} {
		if !strings.Contains(body, want) {
			t.Errorf("the page lacks %q", want)
		}
	}
	// Without a metrics store, nothing is said about GPU time.
	if body := mirrorGet(t, usagePage(t, &fakeScopeSeries{}, nil), "/admin/usage").Body.String(); strings.Contains(body, "GPU time") {
		t.Error("GPU time shown with no metrics store to read it from")
	}
}