89 lines
3.9 KiB
Python
89 lines
3.9 KiB
Python
import http.server
|
|
import json
|
|
import os
|
|
import tempfile
|
|
import threading
|
|
import unittest
|
|
|
|
from werkrevisor.openrouter_activity import ActivityError, calls, import_file, model_name
|
|
from werkrevisor.statsclient import StatsClient
|
|
|
|
CSV = """generation_id,created_at,cost_total,tokens_prompt,tokens_completion,tokens_reasoning,model_permaslug,cancelled
|
|
gen-2,2026-10-06 09:13:33.360,0.067288,29579,700,0,anthropic/claude-sonnet-5.5-20260928,false
|
|
gen-1,2026-10-05 04:10:29.722,0.006916,2134,1446,1200,deepseek/deepseek-v4-pro-20260423,false
|
|
"""
|
|
|
|
|
|
class Service(http.server.BaseHTTPRequestHandler):
|
|
"""Stores each call id once, as the statistics service does."""
|
|
ids = set()
|
|
received = []
|
|
|
|
def do_POST(self):
|
|
body = json.loads(self.rfile.read(int(self.headers["Content-Length"])))
|
|
Service.received.append((self.path, self.headers["Authorization"], body))
|
|
new = [c for c in body["calls"] if c["id"] not in Service.ids]
|
|
Service.ids.update(c["id"] for c in new)
|
|
answer = json.dumps({"stored": len(new), "known": len(body["calls"]) - len(new), "skipped": 0}).encode()
|
|
self.send_response(200)
|
|
self.send_header("Content-Length", str(len(answer)))
|
|
self.end_headers()
|
|
self.wfile.write(answer)
|
|
|
|
def log_message(self, *args):
|
|
pass
|
|
|
|
|
|
class ImportOpenRouterActivityTest(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls.server = http.server.HTTPServer(("127.0.0.1", 0), Service)
|
|
threading.Thread(target=cls.server.serve_forever, daemon=True).start()
|
|
cls.client = StatsClient(f"http://127.0.0.1:{cls.server.server_port}", "source-token")
|
|
fd, cls.csv = tempfile.mkstemp(suffix=".csv")
|
|
with os.fdopen(fd, "w") as f:
|
|
f.write(CSV)
|
|
|
|
@classmethod
|
|
def tearDownClass(cls):
|
|
cls.server.shutdown()
|
|
cls.server.server_close()
|
|
os.remove(cls.csv)
|
|
|
|
def setUp(self):
|
|
Service.ids, Service.received = set(), []
|
|
|
|
def test_sends_the_calls_in_utc_with_the_model_names_of_the_usage_records(self):
|
|
self.assertEqual(import_file(self.csv, self.client), {"stored": 2, "known": 0, "skipped": 0})
|
|
[(path, authorization, body)] = Service.received
|
|
self.assertEqual((path, authorization, body["provider"]), ("/api/billed-calls", "Bearer source-token", "openrouter"))
|
|
self.assertEqual(body["calls"][0], {
|
|
"id": "gen-2", "timestamp": "2026-10-06T09:13:33.360Z", "model": "openrouter/anthropic/claude-sonnet-5.5",
|
|
"cost_usd": 0.067288, "input_tokens": 29579, "output_tokens": 700,
|
|
})
|
|
|
|
def test_a_repeated_import_stores_nothing_new(self):
|
|
import_file(self.csv, self.client)
|
|
self.assertEqual(import_file(self.csv, self.client), {"stored": 0, "known": 2, "skipped": 0})
|
|
|
|
def test_a_cancelled_call_counts_with_its_billed_cost(self):
|
|
[call] = calls([{"generation_id": "gen-3", "created_at": "2026-10-05 04:16:05.057", "cost_total": "0.0173",
|
|
"tokens_prompt": "2011", "tokens_completion": "", "model_permaslug": "openai/gpt-6-astra-20260903",
|
|
"cancelled": "true"}])
|
|
self.assertEqual((call["cost_usd"], call["output_tokens"]), (0.0173, 0))
|
|
|
|
def test_a_malformed_csv_names_the_problem(self):
|
|
for row in ({"generation_id": "gen-4"}, {"generation_id": "", "created_at": "x"},
|
|
{"generation_id": "gen-5", "created_at": "2026-10-05 04:16:05", "cost_total": "free",
|
|
"tokens_prompt": "1", "tokens_completion": "1", "model_permaslug": "m"}):
|
|
with self.subTest(row=row), self.assertRaisesRegex(ActivityError, "malformed activity CSV"):
|
|
calls([row])
|
|
|
|
def test_model_names_lose_only_a_date_suffix(self):
|
|
self.assertEqual(model_name("openai/gpt-6-astra-20260903"), "openrouter/openai/gpt-6-astra")
|
|
self.assertEqual(model_name("deepseek/deepseek-v4-pro"), "openrouter/deepseek/deepseek-v4-pro")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|