diff options
Diffstat (limited to 'internal/storage/dashboard_test.go')
| -rw-r--r-- | internal/storage/dashboard_test.go | 151 |
1 files changed, 151 insertions, 0 deletions
diff --git a/internal/storage/dashboard_test.go b/internal/storage/dashboard_test.go new file mode 100644 index 0000000..b7e425d --- /dev/null +++ b/internal/storage/dashboard_test.go @@ -0,0 +1,151 @@ +package storage + +import ( + "strconv" + "testing" + "time" + + "github.com/thepeterstone/claudomator/internal/task" +) + +// mustCreateTask is a small helper mirroring the pattern used in db_test.go +// (TestCreateTask_AndGetTask) — a minimal valid task row so executions can +// reference a real task_id. +func mustCreateTask(t *testing.T, db *DB, id string) { + t.Helper() + now := time.Now().UTC().Truncate(time.Second) + tk := &task.Task{ + ID: id, + Name: "dashboard-test-" + id, + State: task.StatePending, + DependsOn: []string{}, + CreatedAt: now, + UpdatedAt: now, + } + if err := db.CreateTask(tk); err != nil { + t.Fatalf("CreateTask(%s): %v", id, err) + } +} + +func mustCreateExecution(t *testing.T, db *DB, id, taskID, agent string, rung int, cost float64, start time.Time, status string) { + t.Helper() + exec := &Execution{ + ID: id, + TaskID: taskID, + StartTime: start, + EndTime: start.Add(time.Minute), + Status: status, + CostUSD: cost, + Agent: agent, + EscalationRung: rung, + } + if err := db.CreateExecution(exec); err != nil { + t.Fatalf("CreateExecution(%s): %v", id, err) + } +} + +func TestQueryEscalationFunnel_AggregatesByRungAndAgent(t *testing.T) { + db := testDB(t) + mustCreateTask(t, db, "t1") + + now := time.Now().UTC() + // Rung 0: 2 local executions. + mustCreateExecution(t, db, "e1", "t1", "local", 0, 0, now.Add(-3*time.Hour), "COMPLETED") + mustCreateExecution(t, db, "e2", "t1", "local", 0, 0, now.Add(-2*time.Hour), "COMPLETED") + // Rung 1: 1 anthropic execution. + mustCreateExecution(t, db, "e3", "t1", "anthropic", 1, 0.05, now.Add(-1*time.Hour), "COMPLETED") + // Still in flight — must be excluded. + mustCreateExecution(t, db, "e4", "t1", "anthropic", 1, 0, now, "RUNNING") + // Too old — outside the window. + mustCreateExecution(t, db, "e5", "t1", "local", 0, 0, now.Add(-48*time.Hour), "COMPLETED") + + buckets, err := db.QueryEscalationFunnel(now.Add(-24 * time.Hour)) + if err != nil { + t.Fatalf("QueryEscalationFunnel: %v", err) + } + + want := map[string]struct { + count int + cost float64 + }{ + "0|local": {2, 0}, + "1|anthropic": {1, 0.05}, + } + if len(buckets) != len(want) { + t.Fatalf("expected %d buckets, got %d: %+v", len(want), len(buckets), buckets) + } + for _, b := range buckets { + key := strconv.Itoa(b.Rung) + "|" + b.Agent + w, ok := want[key] + if !ok { + t.Errorf("unexpected bucket %+v", b) + continue + } + if b.Count != w.count { + t.Errorf("bucket %s: count = %d, want %d", key, b.Count, w.count) + } + if b.CostUSD != w.cost { + t.Errorf("bucket %s: cost = %v, want %v", key, b.CostUSD, w.cost) + } + } +} + +func TestQueryEscalationFunnel_Empty_ReturnsEmptySlice(t *testing.T) { + db := testDB(t) + buckets, err := db.QueryEscalationFunnel(time.Now().Add(-24 * time.Hour)) + if err != nil { + t.Fatalf("QueryEscalationFunnel: %v", err) + } + if buckets == nil || len(buckets) != 0 { + t.Errorf("expected empty non-nil slice, got %v", buckets) + } +} + +func TestQuerySpendTimeseries_HourlyBucketsByProvider(t *testing.T) { + db := testDB(t) + mustCreateTask(t, db, "t1") + + // Pin two executions to the exact same hour bucket for two providers. + hourStart := time.Now().UTC().Truncate(time.Hour) + mustCreateExecution(t, db, "e1", "t1", "local", 0, 0.10, hourStart.Add(5*time.Minute), "COMPLETED") + mustCreateExecution(t, db, "e2", "t1", "anthropic", 0, 0.20, hourStart.Add(10*time.Minute), "COMPLETED") + mustCreateExecution(t, db, "e3", "t1", "anthropic", 0, 0.05, hourStart.Add(40*time.Minute), "COMPLETED") + + points, err := db.QuerySpendTimeseries(hourStart.Add(-time.Hour), true) + if err != nil { + t.Fatalf("QuerySpendTimeseries: %v", err) + } + + totals := map[string]float64{} + for _, p := range points { + totals[p.Agent] += p.CostUSD + } + if totals["local"] != 0.10 { + t.Errorf("local total = %v, want 0.10", totals["local"]) + } + if got, want := totals["anthropic"], 0.25; got < want-1e-9 || got > want+1e-9 { + t.Errorf("anthropic total = %v, want %v", got, want) + } +} + +func TestQuerySpendTimeseries_DailyBucketing(t *testing.T) { + db := testDB(t) + mustCreateTask(t, db, "t1") + + today := time.Now().UTC() + mustCreateExecution(t, db, "e1", "t1", "local", 0, 1.5, today, "COMPLETED") + + points, err := db.QuerySpendTimeseries(today.Add(-24*time.Hour), false) + if err != nil { + t.Fatalf("QuerySpendTimeseries: %v", err) + } + if len(points) != 1 { + t.Fatalf("expected 1 point, got %d: %+v", len(points), points) + } + if points[0].Bucket != today.Format("2006-01-02") { + t.Errorf("bucket = %q, want %q", points[0].Bucket, today.Format("2006-01-02")) + } + if points[0].CostUSD != 1.5 { + t.Errorf("cost = %v, want 1.5", points[0].CostUSD) + } +} |
