Repository navigation
Expand file tree
/
Copy pathusage_cli_test.go
More file actions
106 lines (100 loc) · 4.52 KB
/
Copy pathusage_cli_test.go
File metadata and controls
106 lines (100 loc) · 4.52 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
package main
import (
"encoding/csv"
"strings"
"testing"
"time"
"github.com/yetone/magpie/internal/provider"
"github.com/yetone/magpie/internal/sessions"
stats "github.com/yetone/magpie/internal/usage"
)
// magpie usage --csv writes the period's requests, newest first, one row
// each, with the model asked for, sent and served.
func TestUsageCSV(t *testing.T) {
groupsHome(t)
y, m, d := time.Now().Date()
day := time.Date(y, m, d, 0, 0, 0, 0, time.Local) // today's, whenever this runs
stats.Append(stats.Record{Time: day.Add(time.Second), Agent: "codex", Provider: "a", Model: "m", Requested: "fast", Served: "m-mini", Input: 10, Output: 2, Millis: 700, Status: 200})
stats.Append(stats.Record{Time: day.Add(2 * time.Second), Agent: "claude", Provider: "b", Model: "gpt-5.5", Requested: "b/gpt-5.5", Millis: 30, Status: 502})
stats.Append(stats.Record{Time: day.AddDate(0, 0, -3), Agent: "pi", Provider: "a", Model: "m", Input: 1, Status: 200})
var b strings.Builder
if err := usageTo(&b, []string{"usage", "--csv", "today"}); err != nil {
t.Fatal(err)
}
recs, err := csv.NewReader(strings.NewReader(b.String())).ReadAll()
if err != nil {
t.Fatal(err)
}
if len(recs) != 3 || strings.Join(recs[0], ",") != strings.Join(stats.CSVHeader, ",") {
t.Fatalf("%q", recs)
}
col := func(row []string, name string) string {
for i, h := range recs[0] {
if h == name {
return row[i]
}
}
t.Fatalf("no column %s", name)
return ""
}
if r := recs[1]; col(r, "agent") != "claude" || col(r, "status") != "502" || col(r, "error") != "true" || col(r, "requested_model") != "b/gpt-5.5" {
t.Errorf("newest %q", r)
}
if r := recs[2]; col(r, "requested_model") != "fast" || col(r, "model") != "m" || col(r, "served_model") != "m-mini" || col(r, "swapped") != "true" ||
col(r, "input_tokens") != "10" || col(r, "output_tokens") != "2" || col(r, "duration_ms") != "700" || col(r, "error") != "false" {
t.Errorf("swapped %q", r)
}
b.Reset()
if err := usageTo(&b, []string{"usage", "all", "--csv"}); err != nil || strings.Count(b.String(), "\n") != 4 {
t.Fatalf("all: %v %q", err, b.String())
}
if err := usageTo(&b, []string{"usage", "--csv", "today", "extra"}); err == nil {
t.Fatal("took an extra argument")
}
}
func TestUsageProviderKeys(t *testing.T) {
groupsHome(t)
id := provider.KeyID("fixture-provider-secret")
stats.Append(stats.Record{Provider: "relay", Model: "m", ProviderKeyID: id, ProviderKeyName: "Team", Input: 30, Status: 200})
out, err := said(t, func() error { return usageCmd([]string{"usage", "all"}) })
if err != nil || !strings.Contains(out, "upstream provider keys") || !strings.Contains(out, "relay / Team") || strings.Contains(out, "fixture-provider-secret") {
t.Fatal(out, err)
}
var csv strings.Builder
if err := usageTo(&csv, []string{"usage", "--csv", "all"}); err != nil {
t.Fatal(err)
}
if !strings.Contains(csv.String(), "provider_key_id,provider_key_name") || !strings.Contains(csv.String(), id+",Team") || strings.Contains(csv.String(), "fixture-provider-secret") {
t.Fatal(csv.String())
}
}
// magpie usage shows, beside the gateway's calls, those the agents made on
// their own, read from their session files, as the window's Usage page
// does (Kumo31 on Discord: Codex used outside magpie wasn't in it).
func TestUsageShowsCallsNotThroughMagpie(t *testing.T) {
groupsHome(t)
now := time.Now()
stats.Append(stats.Record{Time: now.Add(-time.Minute), Agent: "claude", Provider: "relay", Model: "m", Input: 5, Status: 200})
old := stats.LogCalls
stats.LogCalls = func(time.Time) []sessions.Call {
return []sessions.Call{{Time: now.Add(-2 * time.Minute), Agent: "codex", Session: "c1", Model: "gpt-6-luna", Tokens: sessions.Tokens{Input: 1200, Output: 300}}}
}
t.Cleanup(func() { stats.LogCalls = old })
var b strings.Builder
if err := usageTo(&b, []string{"usage", "today"}); err != nil {
t.Fatal(err)
}
out := b.String()
gw, own, ok := strings.Cut(out, "not through magpie")
if !ok || !strings.Contains(gw, "relay/m") || strings.Contains(gw, "gpt-6-luna") || !strings.Contains(own, "gpt-6-luna") || !strings.Contains(own, "1.5K") {
t.Fatal(out)
}
// a call of the agents' own from before this month shows under all
stats.LogCalls = func(time.Time) []sessions.Call {
return []sessions.Call{{Time: now.AddDate(0, 0, -40), Agent: "codex", Model: "gpt-6-luna", Tokens: sessions.Tokens{Input: 7}}}
}
b.Reset()
if err := usageTo(&b, []string{"usage", "all"}); err != nil || !strings.Contains(b.String(), "not through magpie") || !strings.Contains(b.String(), "gpt-6-luna") {
t.Fatal(b.String(), err)
}
}