Ilker ManapandClaude Opus 5.5
57739d736a
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
49 lines
1.7 KiB
Go
49 lines
1.7 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")
|
|
}
|
|
}
|