diff options
Diffstat (limited to 'tests/test_usage.py')
| -rw-r--r-- | tests/test_usage.py | 337 |
1 files changed, 337 insertions, 0 deletions
diff --git a/tests/test_usage.py b/tests/test_usage.py new file mode 100644 index 0000000..f2d55e3 --- /dev/null +++ b/tests/test_usage.py @@ -0,0 +1,337 @@ +import json +import os +import subprocess +import sys +import tempfile +import unittest +from pathlib import Path + +HERE = Path(__file__).resolve().parent.parent +WF = HERE / "wf.py" +sys.path.insert(0, str(HERE)) + +from wflib import usage # noqa: E402 + + +def entry(mid, req, model, ts, inp=0, cw5=0, cw1h=0, cr=0, out=0, type="assistant", stop=None, content=None, + uuid=None): + u = {"input_tokens": inp, "cache_creation_input_tokens": cw5 + cw1h, "cache_read_input_tokens": cr, + "output_tokens": out} + if cw1h: + u["cache_creation"] = {"ephemeral_5m_input_tokens": cw5, "ephemeral_1h_input_tokens": cw1h} + d = {"type": type, "requestId": req, "timestamp": ts, + "message": {"id": mid, "model": model, "role": "assistant", "usage": u}} + if stop: + d["message"]["stop_reason"] = stop + if content is not None: + d["message"]["content"] = content + if uuid: + d["uuid"] = uuid + return json.dumps(d) + + +# Subagent transcripts: stop_reason always null, output_tokens = message_start placeholder. +# e1: thinking (signature 1800 chars) + text 1000 chars + tool_use (input json 100 chars), placeholder 8. +# by hand: 0.3*(1800-800) + 0.3*1000 + (30 + 0.44*100) = 300 + 300 + 74 = 674 +# e2: same tool_use again but stop_reason set (final): reported 50 is trusted. +# e3: no content, placeholder 9 kept (nothing to estimate). +TOOL = {"type": "tool_use", "id": "t1", "name": "Bash", "input": {"a": "y" * 91}} # json.dumps → 100 chars +SUB = "\n".join([ + entry("e1", "q1", "claude-sonnet-5-5", "2026-10-04T10:00:00Z", out=8, uuid="u1", + content=[{"type": "thinking", "thinking": "", "signature": "z" * 1800}]), + entry("e1", "q1", "claude-sonnet-5-5", "2026-10-04T10:00:01Z", out=8, uuid="u2", + content=[{"type": "text", "text": "x" * 1000}]), + entry("e1", "q1", "claude-sonnet-5-5", "2026-10-04T10:00:01Z", out=8, uuid="u2", + content=[{"type": "text", "text": "x" * 1000}]), # duplicated line: counted once + entry("e1", "q1", "claude-sonnet-5-5", "2026-10-04T10:00:02Z", out=8, uuid="u3", content=[TOOL]), + entry("e2", "q2", "claude-sonnet-5-5", "2026-10-04T10:01:00Z", out=50, stop="tool_use", content=[TOOL]), + entry("e3", "q3", "claude-sonnet-5-5", "2026-10-04T10:02:00Z", out=9), +]) + "\n" + + +# main session, opus 5.5: m1 streamed as 2 entries (same usage), m2 as 2 entries (out 4000 then 10000). +# by hand: m1 = 1000*4 + 100000*5 + 2000*20 = 544000 µ$; m2 = 50000*8 + 1e6*0.2 + 10000*20 = 800000 µ$ +MAIN = "\n".join([ + json.dumps({"type": "user", "timestamp": "2026-10-04T09:00:00Z", "message": {"role": "user", "content": "hi"}}), + entry("m1", "r1", "claude-opus-5-5", "2026-10-04T09:00:01Z", inp=1000, cw5=100000, out=2000), + entry("m1", "r1", "claude-opus-5-5", "2026-10-04T09:00:02Z", inp=1000, cw5=100000, out=2000), + "not json", + json.dumps({"type": "attachment", "timestamp": "2026-10-04T09:00:03Z"}), + entry("m2", "r2", "claude-opus-5-5", "2026-10-04T09:01:00Z", cw1h=50000, cr=1000000, out=4000), + entry("m2", "r2", "claude-opus-5-5", "2026-10-04T09:01:05Z", cw1h=50000, cr=1000000, out=10000), + entry("m9", "r9", "<synthetic>", "2026-10-04T09:02:00Z", out=5), +]) + "\n" + +# subagent, sonnet 5.5: s1 = 10*2 + 20000*2.5 + 30000*0.2 + 500*10 = 61020 µ$; s2 = 50000*0.2 + 1500*10 = 25000 µ$ +AGENT = "\n".join([ + entry("s1", "q1", "claude-sonnet-5-5", "2026-10-04T10:00:00Z", inp=10, cw5=20000, cr=30000, out=500), + entry("s2", "q2", "claude-sonnet-5-5", "2026-10-04T12:00:00Z", cr=50000, out=1500), + entry("s3", "q3", "claude-mystery-1", "2026-10-04T12:30:00Z", inp=7, out=3), +]) + "\n" + +SID = "11111111-2222-3333-4444-555555555555" + + +class Parse(unittest.TestCase): + def test_dedupes_streamed_entries_and_takes_last_output(self): + got = usage.parse(MAIN) + self.assertEqual(list(got), ["claude-opus-5-5"]) + u = got["claude-opus-5-5"] + self.assertEqual((u.turns, u.inp, u.cw5, u.cw1h, u.cr, u.out), (2, 1000, 100000, 50000, 1000000, 12000)) + self.assertEqual(u.cw, 150000) + self.assertAlmostEqual(usage.cost("claude-opus-5-5", u), 1.344) + + def test_models_kept_apart_unknown_price_none(self): + got = usage.parse(AGENT) + s = got["claude-sonnet-5-5"] + self.assertEqual((s.turns, s.cr, s.out), (2, 80000, 2000)) + self.assertAlmostEqual(usage.cost("claude-sonnet-5-5", s), 0.08602) + self.assertIsNone(usage.cost("claude-mystery-1", got["claude-mystery-1"])) + + def test_since_until_filter_on_utc_timestamp(self): + s = usage.parse(AGENT, since="2026-10-04T11:00", until="2026-10-04T12:10")["claude-sonnet-5-5"] + self.assertEqual((s.turns, s.cr), (1, 50000)) + self.assertAlmostEqual(usage.cost("claude-sonnet-5-5", s), 0.025) + + def test_dated_model_id_priced_by_prefix(self): + u = usage.Usage(turns=1, inp=1000000, out=1000000) + self.assertAlmostEqual(usage.cost("claude-haiku-4-5-20251001", u), 6.0) + self.assertAlmostEqual(usage.cost("claude-opus-5", u), 30.0) + + def test_subagent_output_estimated_from_content(self): + u = usage.parse(SUB)["claude-sonnet-5-5"] + self.assertEqual((u.turns, u.out, u.est), (3, 674 + 50 + 9, 1)) + + def test_log_line_marks_estimate(self): + line = usage.log_line("T", "p", "t-x", "1h", "done", "a1", usage.parse(SUB)) + self.assertIn(" out=733 est=1 usd=", line) + self.assertNotIn(" est=", usage.log_line("T", "p", "t-x", "1h", "done", "a1", usage.parse(AGENT))) + + def test_log_line_lane_and_model(self): + by = usage.parse(SUB) + self.assertIn(" lane=fast model=sonnet effort=", usage.log_line("T", "p", "t-x", "1h", "done", "a1", by, lane="fast")) + self.assertIn(" lane=sonnet model=sonnet effort=", usage.log_line("T", "p", "t-x", "1h", "done", "a1", by)) + + def test_report_lane_model_key(self): + es = [{"lane": "fast", "model": "sonnet", "outcome": "done", "usd": 1.0, "turns": 10, "effort": "1h"}, + {"lane": "fast", "model": "opus", "outcome": "done", "usd": 2.0, "turns": 20, "effort": "1h"}, + {"lane": "opus", "outcome": "done", "usd": 3.0, "turns": 30, "effort": "1h"}] + self.assertEqual(usage.report(es, ("1h",)), [ + ("fast/opus", "all", 1, 1, 0, 2.0, 2.0, 2.0, 20, None), ("fast/opus", "1h", 1, 1, 0, 2.0, 2.0, 2.0, 20, None), + ("fast/sonnet", "all", 1, 1, 0, 1.0, 1.0, 1.0, 10, None), ("fast/sonnet", "1h", 1, 1, 0, 1.0, 1.0, 1.0, 10, None), + ("opus", "all", 1, 1, 0, 3.0, 3.0, 3.0, 30, None), ("opus", "1h", 1, 1, 0, 3.0, 3.0, 3.0, 30, None)]) + + def test_fmt(self): + self.assertEqual([usage.fmt(n) for n in (999, 1500, 8948413)], ["999", "1.5k", "8.95M"]) + + +def tcall(mid, ctx, name, inp, tid): + return json.dumps({"type": "assistant", "timestamp": "2026-10-04T10:00:00Z", "message": { + "id": mid, "model": "claude-sonnet-5-5", "usage": {"input_tokens": 10, "cache_read_input_tokens": ctx - 10}, + "content": [{"type": "tool_use", "id": tid, "name": name, "input": inp}]}}) + + +def tres(tid, text): + return json.dumps({"type": "user", "message": {"role": "user", "content": [ + {"type": "tool_result", "tool_use_id": tid, "content": text}]}}) + + +# agent A: ctx 1000 Read(400 chars=100 tok), 2000 grep(800=200), 3000 Edit, 4000 git commit. +# exp 3000 / mut 3000 / book 4000 of 10000; first edit after 2 calls, pre 3000 = 30%, context +2000 +EXA = "\n".join([ + tcall("a1", 1000, "Read", {"file_path": "/w/proj/.worktrees/x/src/a.py"}, "t1"), tres("t1", "x" * 400), + tcall("a2", 2000, "Bash", {"command": "grep -n foo src/a.py 2>/dev/null"}, "t2"), tres("t2", "y" * 800), + tcall("a3", 3000, "Edit", {"file_path": "/w/proj/src/a.py"}, "t3"), tres("t3", "ok"), + tcall("a4", 4000, "Bash", {"command": "git commit -m x"}, "t4"), tres("t4", "done"), +]) + "\n" +# agent B: 4 exploring calls of 1000 each, reads src/a.py (100 tok) + sed -n of b.py (400 chars) +EXB = "\n".join([ + tcall("b1", 1000, "Read", {"file_path": "/w/proj/src/a.py"}, "u1"), tres("u1", "x" * 400), + tcall("b2", 1000, "Bash", {"command": "sed -n 1,9p src/b.py"}, "u2"), tres("u2", "z" * 400), + tcall("b3", 1000, "Bash", {"command": "ls"}, "u3"), tres("u3", "q"), + tcall("b4", 1000, "Glob", {}, "u4"), tres("u4", "q"), +]) + "\n" + + +class Explore(unittest.TestCase): + def test_one_agent(self): + e = usage.explore(EXA, root="/w/proj") + self.assertEqual((e.calls, e.first_edit, e.pre, e.ctx_growth), (4, 2, 3000, 2000)) + self.assertEqual(e.cost, {"exp": 3000, "mut": 3000, "book": 4000}) + self.assertEqual(e.results, {"Read": 100, "grep": 200, "Edit": 0, "other": 1}) + self.assertEqual(e.files, {"src/a.py": 300}) + + def test_call_kind_book_is_writes_only(self): + kind = lambda c: usage.call_kind({"name": "Bash", "input": {"command": c}}) + for c in ("git worktree list", "git branch --list 'fast/*'", "git branch --show-current", "git status --short"): + self.assertEqual(kind(c), "exp", c) + for c in ("git worktree add .w/x -b x master", "git worktree remove .w/x", "git branch -d x", + "python3 /projects/public/workflow/wf.py finish t-x -m ok", "git switch -c x"): + self.assertEqual(kind(c), "book", c) + + def test_since_drops_early_calls(self): + e = usage.explore(EXA.replace("10:00:00Z", "09:00:00Z", 2), since="2026-10-04T10") + self.assertEqual(e.calls, 2) + + def test_report(self): + a, b = usage.explore(EXA, root="/w/proj"), usage.explore(EXB, root="/w/proj") + r = usage.explore_report([("A", a), ("B", b)], min_pre=4) + self.assertEqual(r["rows"], [("A", 4, 2, 0.3, 0.3), ("B", 4, 4, 1.0, 0.0)]) + self.assertEqual((r["total"], r["cost"]["exp"]), (14000, 7000)) + self.assertEqual((r["pre_n"], r["pre_share"], r["pre_calls"], r["pre_ctx"]), (1, 0.3, 2, 2000)) + self.assertEqual(r["results"]["sed-cat"], 100) + self.assertEqual(r["files"], [("src/a.py", 2, 400)]) + self.assertEqual(usage.explore_report([("A", a)], min_calls=5)["rows"], []) + + +class Cli(unittest.TestCase): + def setUp(self): + self.tmp = tempfile.TemporaryDirectory() + self.home = Path(self.tmp.name) + proj = self.home / "projects" / "-work-demo" + sub = proj / SID / "subagents" + sub.mkdir(parents=True) + (proj / f"{SID}.jsonl").write_text(MAIN) + (sub / "agent-abc123.jsonl").write_text(AGENT) + (sub / "agent-abc123.meta.json").write_text(json.dumps({"description": "wf-worker sonnet t-x"})) + + def tearDown(self): + self.tmp.cleanup() + + def wf(self, *args, sid="", cwd=None): + env = {**os.environ, "CLAUDE_CONFIG_DIR": str(self.home), "CLAUDE_CODE_SESSION_ID": sid} + r = subprocess.run([sys.executable, str(WF), "usage", *args], capture_output=True, text=True, + cwd=cwd or self.tmp.name, env=env, timeout=30) + return r.returncode, r.stdout, r.stderr + + def test_explore_skips_main_and_short_agents(self): + code, out, err = self.wf("--session", SID[:8], "--explore") + self.assertEqual((code, err), (0, "")) + self.assertEqual(out.strip(), "no subagent with >= 4 calls") + + def test_session_rows_and_total(self): + code, out, err = self.wf("--session", SID[:8]) + self.assertEqual((code, err), (0, "")) + rows = [l.split() for l in out.strip().split("\n")[1:]] + self.assertEqual(rows[0], ["main", "opus-5-5", "2", "1.0k", "150.0k", "1.00M", "12.0k", "1.34"]) + self.assertEqual(rows[1], ["abc123", "wf-worker", "sonnet", "t-x", + "sonnet-5-5", "2", "10", "20.0k", "80.0k", "2.0k", "0.09"]) + self.assertEqual(rows[2][-7:], ["mystery-1", "1", "7", "0", "0", "3", "?"]) + self.assertEqual(rows[3], ["total", "5", "1.0k", "170.0k", "1.08M", "14.0k", "1.43+"]) + + def test_session_from_env(self): + code, out, _ = self.wf(sid=SID) + self.assertEqual(code, 0) + self.assertIn("abc123", out) + + def test_agent_only(self): + code, out, err = self.wf("--agent", "agent-abc123", "--since", "2026-10-04T11:00") + self.assertEqual((code, err), (0, "")) + lines = out.strip().split("\n") + self.assertNotIn("main", out) + self.assertEqual(lines[1].split()[-7:], ["sonnet-5-5", "1", "0", "0", "50.0k", "1.5k", "0.03"]) + + def test_log_appends_one_key_value_line(self): + root = self.home / "demo" + root.mkdir() + (root / "workflow.toml").write_text("format = 1\n") + code, out, err = self.wf("--agent", "abc123", "--log", "t-x", "done", "--effort", "1h", cwd=root) + self.assertEqual((code, err), (0, "")) + code, out, err = self.wf("--agent", "abc123", "--log", "t-x", "handback", cwd=root) + log = (root / "out" / "wf-cost.log").read_text().split("\n") + self.assertEqual(len(log), 3) # two lines + trailing newline + # lane = model with most turns (sonnet 2 vs mystery 1); usd = known models only + self.assertRegex(log[0], r"^\d{4}-\d\d-\d\dT\d\d:\d\d:\d\dZ project=demo task=t-x lane=sonnet model=sonnet effort=1h " + r"outcome=done turns=3 in=17 cw=20000 cr=80000 out=2003 usd=0\.0860 agent=abc123$") + self.assertIn(" effort=- outcome=handback ", log[1]) + self.assertEqual(out, log[1] + "\n") + + def test_estimated_out_marked_with_tilde(self): + other = self.home / "projects" / "-work-demo" / "other-session" / "subagents" + other.mkdir(parents=True) + (other / "agent-est1.jsonl").write_text(SUB) + code, out, err = self.wf("--agent", "est1") + self.assertEqual((code, err), (0, "")) + self.assertEqual(out.strip().split("\n")[1].split()[-3:], ["0", "~733", "0.01"]) + + def test_log_needs_agent(self): + code, out, err = self.wf("--log", "t-x", "done") + self.assertEqual((code, err), (2, "wf: --log needs --agent\n")) + + def test_log_effort_must_be_an_estimate(self): + code, out, err = self.wf("--agent", "abc123", "--log", "t-x", "done", "--effort", "medium") + self.assertEqual(code, 2) + self.assertIn("invalid choice: 'medium'", err) + + def test_missing_agent_one_line_error(self): + code, out, err = self.wf("--agent", "nope") + self.assertEqual((code, out), (1, "")) + self.assertEqual(err, "wf: no transcript for agent nope\n") + + +LOG_A = """\ +2026-10-01T10:00:00Z project=a task=t-1 lane=sonnet effort=1h outcome=done turns=10 usd=0.50 +2026-10-01T11:00:00Z project=a task=t-2 lane=sonnet effort=1h outcome=handback turns=30 usd=1.50 +garbage line +""" +LOG_B = """\ +2026-10-02T10:00:00Z project=b task=t-3 lane=sonnet effort=<1h outcome=done turns=20 usd=0.40 +2026-10-02T11:00:00Z project=b task=t-4 lane=opus effort=5h outcome=done turns=80 usd=3.00 +2026-10-03T11:00:00Z project=b task=t-5 lane=opus effort=1h outcome=done turns=40 usd=2.00 +""" + + +class Report(unittest.TestCase): + def test_rows_by_lane_then_effort(self): + entries = usage.parse_log(LOG_A + LOG_B) + self.assertEqual(len(entries), 5) + rows = usage.report(entries, ("<1h", "1h", "5h")) + # lane, effort, n, done, other, median $, total $, $/done, median turns — all by hand + self.assertEqual(rows, [ + ("opus", "all", 2, 2, 0, 2.5, 5.0, 2.5, 60, None), + ("opus", "1h", 1, 1, 0, 2.0, 2.0, 2.0, 40, None), + ("opus", "5h", 1, 1, 0, 3.0, 3.0, 3.0, 80, None), + ("sonnet", "all", 3, 2, 1, 0.5, 2.4, 1.2, 20, None), + ("sonnet", "<1h", 1, 1, 0, 0.4, 0.4, 0.4, 20, None), + ("sonnet", "1h", 2, 1, 1, 1.0, 2.0, 2.0, 20, None), + ]) + + def test_duration(self): + line = usage.log_line("T", "p", "t-x", "1h", "done", "a1", usage.parse(SUB), dur=125) + self.assertIn(" dur=125 agent=a1", line) + self.assertNotIn("dur=", usage.log_line("T", "p", "t-x", "1h", "done", "a1", usage.parse(SUB))) + log = ("2026-10-04T10:00:00Z project=a task=t-6 lane=fast effort=<1h outcome=done turns=1 usd=1 dur=100\n" + "2026-10-04T10:00:00Z project=a task=t-7 lane=fast effort=<1h outcome=done turns=1 usd=1 dur=300\n" + "2026-10-04T10:00:00Z project=a task=t-8 lane=fast effort=<1h outcome=done turns=1 usd=1\n") + self.assertEqual(usage.report(usage.parse_log(log), ("<1h",))[0][9], 200) + + def test_done_gate_red_counts_as_done(self): + log = ("2026-10-04T10:00:00Z project=a task=t-6 lane=fast effort=<1h outcome=done+gate-red turns=12 usd=1.00\n" + "2026-10-04T11:00:00Z project=a task=t-7 lane=fast effort=<1h outcome=donex turns=8 usd=0.50\n") + rows = usage.report(usage.parse_log(log), ("<1h",)) + self.assertEqual(rows[0], ("fast", "all", 2, 1, 1, 0.75, 1.5, 1.5, 10, None)) + + def test_since_and_no_done(self): + entries = usage.parse_log(LOG_A + LOG_B, since="2026-10-01T10:30") + rows = usage.report(entries, ("1h",)) + self.assertIn(("sonnet", "all", 2, 1, 1, 0.95, 1.9, 1.9, 25, None), rows) + rows = usage.report(usage.parse_log(LOG_A, since="2026-10-01T10:30"), ("1h",)) + self.assertEqual(rows[0], ("sonnet", "all", 1, 0, 1, 1.5, 1.5, None, 30, None)) + + def test_cli_reads_every_project(self): + with tempfile.TemporaryDirectory() as tmp: + for name, text in (("a", LOG_A), ("b", LOG_B)): + (Path(tmp) / name / "out").mkdir(parents=True) + (Path(tmp) / name / "workflow.toml").write_text("format = 1\n") + (Path(tmp) / name / "out" / "wf-cost.log").write_text(text) + r = subprocess.run([sys.executable, str(WF), "usage", "--report"], capture_output=True, text=True, + cwd=tmp, env={**os.environ, "WF_ROOT": tmp}, timeout=30) + self.assertEqual((r.returncode, r.stderr), (0, "")) + lines = [l.split() for l in r.stdout.strip().split("\n")] + self.assertEqual(lines[0], ["lane", "effort", "n", "done", "other", "med$", "total$", "$/done", "med_turns", "med_dur"]) + self.assertEqual(lines[1], ["opus", "all", "2", "2", "0", "2.50", "5.00", "2.50", "60", "-"]) + self.assertEqual(lines[6], ["sonnet", "1h", "2", "1", "1", "1.00", "2.00", "2.00", "20", "-"]) + + +if __name__ == "__main__": + unittest.main() |