eleven / ocho Public

Download this file

1	package main
2
3	import (
4	"bytes"
5	"encoding/json"
6	"net/http"
7	"net/http/httptest"
8	"testing"
9	"time"
10	)
11
12	func ingestUsage(t *testing.T, s *cloudServer, token string, rows []botUsageRow) (int, map[string]any) {
13	t.Helper()
14	return ingestUsageReport(t, s, token, "bbbbbbbbbbbbbbbb", rows)
15	}
16
17	func ingestUsageReport(t *testing.T, s *cloudServer, token, reportID string, rows []botUsageRow) (int, map[string]any) {
18	t.Helper()
19	raw, _ := json.Marshal(map[string]any{"report_id": reportID, "rows": rows})
20	req := httptest.NewRequest("POST", "/api/usage", bytes.NewReader(raw))
21	if token != "" {
22	req.Header.Set("Authorization", "Bearer "+token)
23	}
24	rec := httptest.NewRecorder()
25	s.usageHandler(rec, req)
26	var body map[string]any
27	json.NewDecoder(rec.Body).Decode(&body)
28	return rec.Code, body
29	}
30
31	// Ingest: the meter token resolves to exactly one space, rows upsert
32	// additively (partial-day reports accumulate), and a stranger is refused.
33	func TestUsageIngest(t *testing.T) {
34	db := openTestCloudDB(t)
35	id := insertSpaceRow(t, db, "space.11ch.at", 1000)
36	if err := db.setMeterTokenHash(id, hashToken("meter-tok")); err != nil {
37	t.Fatal(err)
38	}
39	s := &cloudServer{db: db}
40
41	rows := []botUsageRow{
42	{Day: "2026-07-14", Model: "claude-sonnet-5", Dimension: "llm_input_tokens", Quantity: 100},
43	{Day: "2026-07-14", Dimension: "charter_runs", Quantity: 2},
44	}
45	if code, body := ingestUsageReport(t, s, "meter-tok", "aa11aa11aa11aa11", rows); code != http.StatusOK {
46	t.Fatalf("ingest: %d %v", code, body)
47	}
48	// A second partial-day report (new id) adds, never clobbers.
49	if code, _ := ingestUsageReport(t, s, "meter-tok", "aa22aa22aa22aa22", rows[:1]); code != http.StatusOK {
50	t.Fatal("second ingest failed")
51	}
52	// A REPLAY (same id — the ack-lost retry) is recognized and dropped whole.
53	if code, _ := ingestUsageReport(t, s, "meter-tok", "aa11aa11aa11aa11", rows); code != http.StatusOK {
54	t.Fatal("replay must still answer ok")
55	}
56	got, err := db.usageFor(id)
57	if err != nil {
58	t.Fatal(err)
59	}
60	byKey := map[string]int64{}
61	for _, r := range got {
62	byKey[r.Model+"|"+r.Dimension] = r.Quantity
63	}
64	if byKey["claude-sonnet-5|llm_input_tokens"] != 200 || byKey["|charter_runs"] != 2 {
65	t.Fatalf("ledger = %v (replay must not double-count)", byKey)
66	}
67	// Prune drops old dedupe ids; a post-prune replay would re-count, which
68	// is why the window is days against a minutes-scale retry horizon.
69	if err := db.pruneUsageReports(time.Now().Unix() + 1); err != nil {
70	t.Fatal(err)
71	}
72
73	if code, _ := ingestUsage(t, s, "wrong", rows); code != http.StatusForbidden {
74	t.Fatal("bad token must be refused")
75	}
76	if code, _ := ingestUsage(t, s, "", rows); code != http.StatusForbidden {
77	t.Fatal("missing token must be refused")
78	}
79	// An unmetered space's empty hash must never match any presented token.
80	other := insertSpaceRow(t, db, "other.11ch.at", 1000)
81	if sid, ok := db.spaceByMeterToken(""); ok {
82	t.Fatalf("empty token resolved to space %d", sid)
83	}
84	_ = other
85	if code, _ := ingestUsage(t, s, "meter-tok", []botUsageRow{{Day: "bad", Dimension: "llm_requests", Quantity: 1}}); code != http.StatusBadRequest {
86	t.Fatal("invalid rows must be refused at the ledger too")
87	}
88	}
89
90	// The sweep mints a meter token for any space without one, pushes it through
91	// controld, and records only the hash. A space that already has one is left
92	// alone — tokens rotate only when a push failed to record.
93	func TestMeterSweepMintsTokens(t *testing.T) {
94	db := openTestCloudDB(t)
95	id := insertSpaceRow(t, db, "fresh.11ch.at", 1000)
96
97	var pushes int
98	var gotURL, gotToken string
99	mux := http.NewServeMux()
100	mux.HandleFunc("GET /instances/{host}/stats", func(w http.ResponseWriter, r *http.Request) {
101	writeJSON(w, http.StatusOK, map[string]any{"storage_bytes": 1, "setup_complete": true, "members": 1})
102	})
103	mux.HandleFunc("POST /instances/{host}/meter", func(w http.ResponseWriter, r *http.Request) {
104	var body struct{ URL, Token string }
105	json.NewDecoder(r.Body).Decode(&body)
106	pushes++
107	gotURL, gotToken = body.URL, body.Token
108	writeJSON(w, http.StatusOK, map[string]any{"ok": true})
109	})
110	sock := fakeControl(t, mux.ServeHTTP)
111	s := &cloudServer{db: db, control: controlClient(sock), controlToken: "t",
112	usageURL: "https://cloud.example/api/usage"}
113
114	s.meterOnce()
115	if pushes != 1 || gotURL != "https://cloud.example/api/usage" || gotToken == "" {
116	t.Fatalf("push: n=%d url=%q token-empty=%v", pushes, gotURL, gotToken == "")
117	}
118	// The pushed token authenticates; only its hash is stored.
119	if sid, ok := db.spaceByMeterToken(gotToken); !ok || sid != id {
120	t.Fatalf("pushed token does not resolve: ok=%v sid=%d", ok, sid)
121	}
122	// Converged: the next sweep must not re-mint.
123	s.meterOnce()
124	if pushes != 1 {
125	t.Fatalf("second sweep re-minted (pushes=%d)", pushes)
126	}
127	}
128
129	// With the usage pipe off (no LCHAT_CLOUD_USAGE_URL) the sweep never touches
130	// meter tokens — self-hosted cloud deployments see no behavior change.
131	func TestMeterSweepNoUsageURL(t *testing.T) {
132	db := openTestCloudDB(t)
133	insertSpaceRow(t, db, "fresh.11ch.at", 1000)
134	mux := http.NewServeMux()
135	mux.HandleFunc("GET /instances/{host}/stats", func(w http.ResponseWriter, r *http.Request) {
136	writeJSON(w, http.StatusOK, map[string]any{"storage_bytes": 1})
137	})
138	mux.HandleFunc("POST /instances/{host}/meter", func(w http.ResponseWriter, r *http.Request) {
139	t.Error("meter push must not happen with the pipe off")
140	})
141	sock := fakeControl(t, mux.ServeHTTP)
142	s := &cloudServer{db: db, control: controlClient(sock), controlToken: "t"}
143	s.meterOnce()
144	}
145