llmgw/internal/store/export_test.go

package store

import (
	"testing"
	"time"
)

func TestUsageIsExportedByPeriodAndWhatWasAskedFor(t *testing.T) {
	db := capacityStore(t)
	ctx := t.Context()
	org, _ := db.CreateOrg(ctx, "acme")
	team, _ := db.CreateTeam(ctx, org, "ml")
	project, _ := db.CreateProject(ctx, team, "train")
	day1 := time.Date(2026, 9, 30, 10, 0, 0, 0, time.UTC)
	day2 := time.Date(2026, 10, 1, 9, 0, 0, 0, time.UTC)
	var rows []Usage
	for _, r := range []struct {
		at    time.Time
		model string
		p     int64
	}{{day1, "a", 10}, {day1, "b", 20}, {day2, "a", 5}, {day2.Add(time.Hour), "a", 5}} {
		rows = append(rows, Usage{At: r.at, KeyID: "k", ProjectID: project, Model: r.model, ServedModel: r.model,
			PromptTokens: r.p, CompletionTokens: 1, CostMicros: 100, Status: 200})
	}
	if err := db.WriteUsage(ctx, rows); err != nil {
		t.Fatal(err)
	}
	from, to := day1.Add(-24*time.Hour), day2.Add(48*time.Hour)

	byDay, err := db.UsageExport(ctx, from, to, "day", []string{"project", "model"})
	if err != nil || len(byDay) != 3 {
		t.Fatalf("%+v %v", byDay, err)
	}
	last := byDay[2]
	if last.Project != "train" || last.Model != "a" || last.Prompt != 10 || last.Requests != 2 || !last.Period.Equal(time.Date(2026, 10, 1, 0, 0, 0, 0, time.UTC)) {
		t.Errorf("%+v", last)
	}
	byMonth, _ := db.UsageExport(ctx, from, to, "month", []string{"org"})
	if len(byMonth) != 2 || byMonth[0].Org != "acme" || byMonth[0].Prompt != 30 || byMonth[1].CostMicros != 200 {
		t.Errorf("by month: %+v", byMonth)
	}
	if _, err := db.UsageExport(ctx, from, to, "week", nil); err == nil {
		t.Error("a week was taken")
	}
	if _, err := db.UsageExport(ctx, from, to, "day", []string{"planet"}); err == nil {
		t.Error("an unknown dimension was taken")
	}
}