| 1 | package main |
| 2 | |
| 3 | import ( |
| 4 | "net/http" |
| 5 | "testing" |
| 6 | "time" |
| 7 | ) |
| 8 | |
| 9 | func testRates() map[string]aiRates { |
| 10 | return map[string]aiRates{ |
| 11 | // Round numbers so the arithmetic below reads at a glance: |
| 12 | // €10/M in, €50/M out, €12.50/M cache-write, €1/M cache-read. |
| 13 | "claude-sonnet-5": {In: 10, Out: 50, CacheWrite: 12.5, CacheRead: 1}, |
| 14 | } |
| 15 | } |
| 16 | |
| 17 | func TestParseAIRates(t *testing.T) { |
| 18 | rates, err := parseAIRates(`{"claude-sonnet-5":{"in":3,"out":15,"cw":3.75,"cr":0.3}}`) |
| 19 | if err != nil || rates["claude-sonnet-5"].Out != 15 { |
| 20 | t.Fatalf("parse: %v %+v", err, rates) |
| 21 | } |
| 22 | if r, err := parseAIRates(""); err != nil || r != nil { |
| 23 | t.Fatalf("empty must be off: %v %v", r, err) |
| 24 | } |
| 25 | if _, err := parseAIRates(`{"m":{"in":-1}}`); err == nil { |
| 26 | t.Fatal("negative rates must be refused") |
| 27 | } |
| 28 | if _, err := parseAIRates(`nonsense`); err == nil { |
| 29 | t.Fatal("malformed JSON must be refused") |
| 30 | } |
| 31 | } |
| 32 | |
| 33 | func TestUsageCostEUR(t *testing.T) { |
| 34 | rows := []botUsageRow{ |
| 35 | {Model: "claude-sonnet-5", Dimension: "llm_input_tokens", Quantity: 1_000_000}, // €10 |
| 36 | {Model: "claude-sonnet-5", Dimension: "llm_output_tokens", Quantity: 200_000}, // €10 |
| 37 | {Model: "claude-sonnet-5", Dimension: "llm_cache_read_tokens", Quantity: 3_000_000}, // €3 |
| 38 | {Model: "claude-sonnet-5", Dimension: "llm_requests", Quantity: 42}, // request counts are free |
| 39 | {Dimension: "charter_runs", Quantity: 7}, // activity, not money |
| 40 | {Model: "unpriced-model", Dimension: "llm_output_tokens", Quantity: 1_000_000}, // no rate → 0 + a log line |
| 41 | } |
| 42 | if got := usageCostEUR(rows, testRates()); got != 23 { |
| 43 | t.Fatalf("cost = %v, want 23", got) |
| 44 | } |
| 45 | } |
| 46 | |
| 47 | // The state machine end to end through the ingest handler: under allowance is |
| 48 | // ok, over allowance is exhausted, and a drained fleet pool exhausts a space |
| 49 | // that is itself under its allowance. |
| 50 | func TestBudgetStateThroughIngest(t *testing.T) { |
| 51 | db := openTestCloudDB(t) |
| 52 | a := insertSpaceRow(t, db, "a.11ch.at", 1000) |
| 53 | b := insertSpaceRow(t, db, "b.11ch.at", 1000) |
| 54 | for id, tok := range map[int64]string{a: "tok-a", b: "tok-b"} { |
| 55 | if err := db.setMeterTokenHash(id, hashToken(tok)); err != nil { |
| 56 | t.Fatal(err) |
| 57 | } |
| 58 | } |
| 59 | s := &cloudServer{db: db, aiRates: testRates(), aiAllowanceEUR: 5, aiPoolEUR: 12} |
| 60 | day := time.Now().UTC().Format("2006-01-02") |
| 61 | |
| 62 | report := func(tok, reportID string, outTokens int64) map[string]any { |
| 63 | t.Helper() |
| 64 | code, body := ingestUsageReport(t, s, tok, reportID, []botUsageRow{ |
| 65 | {Day: day, Model: "claude-sonnet-5", Dimension: "llm_output_tokens", Quantity: outTokens}}) |
| 66 | if code != http.StatusOK { |
| 67 | t.Fatalf("ingest: %d %v", code, body) |
| 68 | } |
| 69 | return body["budget"].(map[string]any) |
| 70 | } |
| 71 | |
| 72 | // €2 of output: comfortably inside the €5 allowance. |
| 73 | if st := report("tok-a", "aaaaaaaaaaaaaaa1", 40_000); st["state"] != "ok" { |
| 74 | t.Fatalf("under allowance: %v", st) |
| 75 | } |
| 76 | // Another €4 → €6 total: over the per-space allowance. |
| 77 | if st := report("tok-a", "aaaaaaaaaaaaaaa2", 80_000); st["state"] != "exhausted" { |
| 78 | t.Fatalf("over allowance: %v", st) |
| 79 | } |
| 80 | // Space B is fresh (€4 < €5 allowance)… but the fleet pool (€12) is now |
| 81 | // €6 + €4 = €10; one more €4 pushes past it and B exhausts too. |
| 82 | if st := report("tok-b", "bbbbbbbbbbbbbbb1", 80_000); st["state"] != "ok" { |
| 83 | t.Fatalf("pool not yet drained: %v", st) |
| 84 | } |
| 85 | if st := report("tok-b", "bbbbbbbbbbbbbbb2", 80_000); st["state"] != "exhausted" { |
| 86 | t.Fatalf("pool drained must exhaust: %v", st) |
| 87 | } |
| 88 | |
| 89 | // Unconfigured stays measure-only whatever the ledger says. |
| 90 | s2 := &cloudServer{db: db} |
| 91 | if st := s2.budgetState(a); st != "ok" { |
| 92 | t.Fatalf("unconfigured: %q", st) |
| 93 | } |
| 94 | } |
| 95 | |
amadan