Ilker ManapandClaude Opus 5.5
eaea671f23
The five-minute, hourly and daily sums gain images (migration 0042), and the CSV and JSON exports a column for them: once request rows were past their retention, an export said what pictures cost but not how many there were. The pictures made before are added to the sums from their rows. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
82 lines
3.0 KiB
Go
82 lines
3.0 KiB
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")
|
|
}
|
|
}
|
|
|
|
func TestPicturesAreCountedInTheSumsAndTheExport(t *testing.T) {
|
|
// The request rows said how many pictures each made, and the sums did
|
|
// not: once the rows were past their retention, an export said what the
|
|
// pictures cost and not how many there were.
|
|
db := capacityStore(t)
|
|
ctx := t.Context()
|
|
at := time.Date(2026, 10, 10, 9, 0, 0, 0, time.UTC)
|
|
if err := db.WriteUsage(ctx, []Usage{
|
|
{At: at, KeyID: "k", Model: "flux", ServedModel: "flux", PromptTokens: 9, Images: 2, CostMicros: 40000, Status: 200},
|
|
{At: at.Add(time.Minute), KeyID: "k", Model: "flux", ServedModel: "flux", PromptTokens: 9, Images: 3, CostMicros: 60000, Status: 200},
|
|
{At: at.Add(2 * time.Minute), KeyID: "k", Model: "chat", ServedModel: "chat", PromptTokens: 9, Status: 200},
|
|
}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
// Past their retention: only the sums are left.
|
|
if _, err := db.PruneUsage(ctx, at.Add(time.Hour)); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
for _, period := range []string{"hour", "day"} {
|
|
rows, err := db.UsageExport(ctx, at.Add(-24*time.Hour), at.Add(48*time.Hour), period, []string{"model"})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
got := map[string]int64{}
|
|
for _, r := range rows {
|
|
got[r.Model] += r.Images
|
|
}
|
|
if got["flux"] != 5 || got["chat"] != 0 {
|
|
t.Errorf("%s: pictures by model %v, want flux 5", period, got)
|
|
}
|
|
}
|
|
}
|