llmgw/internal/store/usageseries_test.go

package store

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

// spreadOverTime is one org, two teams, and traffic at two separate hours, so a
// bucketed report has something to put in more than one slice.
func spreadOverTime(t *testing.T) (*DB, time.Time) {
	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"}))
	must(s.CreateUser(ctx, User{ID: "usr-1", Email: "ada@example.com", CreatedAt: time.Now()}))
	for id, proj := range map[string]string{"k1": "p1", "k2": "p2"} {
		_, k, _ := NewKey(id, "n", proj)
		k.CreatedAt = time.Now()
		must(s.CreateKey(ctx, k))
	}

	now := time.Now()
	old := now.Add(-3 * time.Hour)
	must(s.WriteUsage(ctx, []Usage{
		// platform, three hours ago
		{At: old, KeyID: "k1", ProjectID: "p1", Model: "m", ServedModel: "m",
			PromptTokens: 100, CompletionTokens: 50, CostMicros: 900, Status: 200},
		// platform, now
		{At: now, KeyID: "k1", ProjectID: "p1", Model: "m", ServedModel: "m",
			PromptTokens: 10, CompletionTokens: 5, CostMicros: 100, Status: 200},
		// research, now only
		{At: now, KeyID: "k2", ProjectID: "p2", Model: "m", ServedModel: "m",
			PromptTokens: 1, CompletionTokens: 1, CostMicros: 50, Status: 500},
	}))
	return s, now
}

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

func TestUsageIsBucketedPerGroup(t *testing.T) {
	s, _ := spreadOverTime(t)

	rows, err := s.UsageBucketsBy(t.Context(), GroupByTeam,
		time.Now().Add(-6*time.Hour), time.Hour)
	if err != nil {
		t.Fatal(err)
	}
	if len(rows) != 2 {
		t.Fatalf("expected two teams, got %d", len(rows))
	}

	platform, ok := seriesFor(rows, "t1")
	if !ok {
		t.Fatal("platform is missing")
	}
	// The totals are the window's, not one bucket's: a table beside the chart
	// reads these rather than adding the buckets up itself.
	if platform.CostMicros != 1000 || platform.Requests != 2 {
		t.Errorf("platform totals wrong: %d micros over %d requests",
			platform.CostMicros, platform.Requests)
	}

	// Two of the seven hourly slices carry traffic, and they are not adjacent.
	busy := 0
	for _, b := range platform.Buckets {
		if b.Requests > 0 {
			busy++
		}
	}
	if busy != 2 {
		t.Errorf("expected traffic in two buckets, found %d", busy)
	}
}

func TestEveryGroupGetsTheSameTimestamps(t *testing.T) {
	// Gap-filled, and the same fill for everyone. A chart sums the tail of its
	// series by timestamp, so two groups whose points do not line up would add
	// up to a line that is wrong everywhere they disagree.
	s, _ := spreadOverTime(t)

	rows, err := s.UsageBucketsBy(t.Context(), GroupByTeam,
		time.Now().Add(-6*time.Hour), time.Hour)
	if err != nil {
		t.Fatal(err)
	}
	if len(rows) < 2 {
		t.Fatalf("expected two teams, got %d", len(rows))
	}
	a, b := rows[0].Buckets, rows[1].Buckets
	if len(a) != len(b) {
		t.Fatalf("one group has %d buckets, the other %d", len(a), len(b))
	}
	for i := range a {
		if !a[i].Start.Equal(b[i].Start) {
			t.Fatalf("bucket %d is %s for one group and %s for the other",
				i, a[i].Start, b[i].Start)
		}
	}
}

func TestAGroupThatWentQuietReadsAsZeroRatherThanAbsent(t *testing.T) {
	// research used the cluster once, in the last hour. The earlier slices must
	// be present and zero: a line drawn straight from three hours ago to now
	// would claim it was busy the whole time.
	s, _ := spreadOverTime(t)

	rows, err := s.UsageBucketsBy(t.Context(), GroupByTeam,
		time.Now().Add(-6*time.Hour), time.Hour)
	if err != nil {
		t.Fatal(err)
	}
	research, ok := seriesFor(rows, "t2")
	if !ok {
		t.Fatal("research is missing")
	}
	if len(research.Buckets) < 6 {
		t.Fatalf("expected the window filled in, got %d buckets", len(research.Buckets))
	}
	if research.Buckets[0].Requests != 0 {
		t.Errorf("the first slice should be an explicit zero, got %d",
			research.Buckets[0].Requests)
	}
	if research.Requests != 1 || research.Errors != 1 {
		t.Errorf("research totals wrong: %d requests, %d errors",
			research.Requests, research.Errors)
	}
}

func TestUsageCanBeGroupedByPerson(t *testing.T) {
	// A person is not a level of the tenancy tree — they hold keys across
	// projects — which is why this grouping exists beside the other three.
	s, _ := spreadOverTime(t)

	rows, err := s.UsageBucketsBy(t.Context(), GroupByUser,
		time.Now().Add(-6*time.Hour), time.Hour)
	if err != nil {
		t.Fatal(err)
	}
	// Both keys were issued without an owner, so the traffic groups under the
	// empty person rather than disappearing.
	unowned, ok := seriesFor(rows, "")
	if !ok {
		t.Fatalf("keys with no owner were dropped: %+v", rows)
	}
	if unowned.Requests != 3 {
		t.Errorf("expected all three requests, got %d", unowned.Requests)
	}
}

func TestAnUnknownGroupingIsRefusedRatherThanBuiltIntoSQL(t *testing.T) {
	s, _ := spreadOverTime(t)

	_, err := s.UsageBucketsBy(t.Context(), Grouping("key_id; DROP TABLE usage"),
		time.Now().Add(-time.Hour), time.Hour)
	if err == nil {
		t.Fatal("an unknown grouping was accepted")
	}
}