from __future__ import annotations import json import sqlite3 import tempfile import unittest from datetime import datetime from pathlib import Path from zoneinfo import ZoneInfo from codex_usage_profiler.core import build_report from codex_usage_profiler.render import render_html, render_markdown class UsageProfilerTests(unittest.TestCase): def test_aggregates_threads_edges_and_rollout_breakdowns(self) -> None: with tempfile.TemporaryDirectory() as tmp: root = Path(tmp) state = root / "state_5.sqlite" sessions_dir = root / "sessions" sessions_dir.mkdir() main_rollout = sessions_dir / "rollout-main.jsonl" child_rollout = sessions_dir / "rollout-child.jsonl" _write_rollout( main_rollout, [ {"type": "turn_context", "payload": {"model": "gpt-5.4", "cwd": "/home/test/project"}}, {"type": "event_msg", "payload": {"type": "task_started"}}, { "type": "event_msg", "payload": { "type": "token_count", "info": { "total_token_usage": { "input_tokens": 70, "cached_input_tokens": 20, "output_tokens": 25, "reasoning_output_tokens": 5, "total_tokens": 100, } }, }, }, ], malformed=True, ) _write_rollout( child_rollout, [ {"type": "turn_context", "payload": {"model": "gpt-5.4-mini", "cwd": "/home/test/project"}}, { "type": "event_msg", "payload": { "type": "token_count", "info": { "total_token_usage": { "input_tokens": 35, "cached_input_tokens": 10, "output_tokens": 12, "reasoning_output_tokens": 3, "total_tokens": 50, } }, }, }, ], ) _create_state( state, [ { "id": "main", "rollout_path": str(main_rollout), "created_at_ms": _epoch_ms("2026-04-20T22:30:00+00:00"), "updated_at_ms": _epoch_ms("2026-04-20T23:00:00+00:00"), "model_provider": "openai", "cwd": "/home/test/project", "title": "Main work", "tokens_used": 100, "agent_nickname": "", "agent_role": "", "model": None, }, { "id": "child", "rollout_path": str(child_rollout), "created_at_ms": _epoch_ms("2026-04-21T10:00:00+00:00"), "updated_at_ms": _epoch_ms("2026-04-21T10:05:00+00:00"), "model_provider": "openai", "cwd": "/home/test/project", "title": "Explore", "tokens_used": 50, "agent_nickname": "Noether", "agent_role": "explorer", "model": "gpt-5.4-mini", }, ], [("main", "child")], ) report = build_report( state_path=state, sessions_dir=sessions_dir, tz_name="Europe/Berlin", now=datetime(2026, 4, 22, 12, tzinfo=ZoneInfo("Europe/Berlin")), ) self.assertEqual(report.total_tokens, 150) self.assertEqual(report.daily[0].label, "2026-04-21") self.assertEqual(report.daily[0].tokens, 150) self.assertEqual(report.weekly[0].label, "2026-W17") self.assertEqual({bucket.label: bucket.tokens for bucket in report.by_kind}, {"main": 100, "subagent": 50}) self.assertEqual({bucket.label: bucket.tokens for bucket in report.by_role}, {"main": 100, "explorer": 50}) self.assertEqual({bucket.label: bucket.tokens for bucket in report.by_model}, {"gpt-5.4": 100, "gpt-5.4-mini": 50}) self.assertEqual(report.rollout_totals.total_tokens, 150) self.assertEqual(report.rollout_totals.cached_input_tokens, 30) self.assertEqual(len(report.warnings), 1) def test_days_filter_uses_local_calendar_days(self) -> None: with tempfile.TemporaryDirectory() as tmp: root = Path(tmp) state = root / "state_5.sqlite" sessions_dir = root / "sessions" sessions_dir.mkdir() _create_state( state, [ { "id": "old", "rollout_path": "", "created_at_ms": _epoch_ms("2026-04-18T12:00:00+00:00"), "updated_at_ms": _epoch_ms("2026-04-18T12:00:00+00:00"), "model_provider": "openai", "cwd": "/tmp/a", "title": "Old", "tokens_used": 99, "agent_nickname": "", "agent_role": "", "model": "gpt-5.4", }, { "id": "new", "rollout_path": "", "created_at_ms": _epoch_ms("2026-04-21T12:00:00+00:00"), "updated_at_ms": _epoch_ms("2026-04-21T12:00:00+00:00"), "model_provider": "openai", "cwd": "/tmp/a", "title": "New", "tokens_used": 10, "agent_nickname": "", "agent_role": "", "model": "gpt-5.4", }, ], [], ) report = build_report( state_path=state, sessions_dir=sessions_dir, tz_name="Europe/Berlin", days=2, now=datetime(2026, 4, 22, 12, tzinfo=ZoneInfo("Europe/Berlin")), ) self.assertEqual(len(report.sessions), 1) self.assertEqual(report.sessions[0].thread_id, "new") self.assertEqual(report.total_tokens, 10) def test_renderers_include_expected_sections_and_escape_html(self) -> None: with tempfile.TemporaryDirectory() as tmp: root = Path(tmp) state = root / "state_5.sqlite" sessions_dir = root / "sessions" sessions_dir.mkdir() _create_state( state, [ { "id": "s1", "rollout_path": "", "created_at_ms": _epoch_ms("2026-04-21T12:00:00+00:00"), "updated_at_ms": _epoch_ms("2026-04-21T12:00:00+00:00"), "model_provider": "openai", "cwd": "/tmp/a", "title": "", "tokens_used": 10, "agent_nickname": "", "agent_role": "", "model": "gpt-5.4", } ], [], ) report = build_report( state_path=state, sessions_dir=sessions_dir, tz_name="Europe/Berlin", now=datetime(2026, 4, 22, 12, tzinfo=ZoneInfo("Europe/Berlin")), ) markdown = render_markdown(report) html = render_html(report) self.assertIn("## Daily Tokens", markdown) self.assertIn("## Top Sessions", markdown) self.assertNotIn("<unsafe>", html) self.assertNotIn("", html) def _create_state(state: Path, rows: list[dict[str, object]], edges: list[tuple[str, str]]) -> None: with sqlite3.connect(state) as conn: conn.execute( """ CREATE TABLE threads ( id TEXT PRIMARY KEY, rollout_path TEXT NOT NULL, created_at INTEGER NOT NULL, updated_at INTEGER NOT NULL, created_at_ms INTEGER, updated_at_ms INTEGER, model_provider TEXT NOT NULL, cwd TEXT NOT NULL, title TEXT NOT NULL, tokens_used INTEGER NOT NULL, agent_nickname TEXT, agent_role TEXT, model TEXT ) """ ) conn.execute( """ CREATE TABLE thread_spawn_edges ( parent_thread_id TEXT NOT NULL, child_thread_id TEXT NOT NULL PRIMARY KEY ) """ ) for row in rows: conn.execute( """ INSERT INTO threads ( id, rollout_path, created_at, updated_at, created_at_ms, updated_at_ms, model_provider, cwd, title, tokens_used, agent_nickname, agent_role, model ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( row["id"], row["rollout_path"], int(row["created_at_ms"]) // 1000, int(row["updated_at_ms"]) // 1000, row["created_at_ms"], row["updated_at_ms"], row["model_provider"], row["cwd"], row["title"], row["tokens_used"], row["agent_nickname"], row["agent_role"], row["model"], ), ) for parent, child in edges: conn.execute( "INSERT INTO thread_spawn_edges (parent_thread_id, child_thread_id) VALUES (?, ?)", (parent, child), ) def _write_rollout(path: Path, events: list[dict[str, object]], malformed: bool = False) -> None: with path.open("w", encoding="utf-8") as handle: for event in events: handle.write(json.dumps(event) + "\n") if malformed: handle.write("{not-json\n") def _epoch_ms(value: str) -> int: return int(datetime.fromisoformat(value).timestamp() * 1000) if __name__ == "__main__": unittest.main()