fuzzyworld-tools / files / tests / test_usage.py
test_usage.py
ced0eafa 2026-11-04 Pavel Safronov | 1 | """Tests for lib/usage.py --trailers. Run: python3 -m unittest discover tests""" |
| 2 | import io | |
| 3 | import os | |
| 4 | import shutil | |
| 5 | import subprocess | |
| 6 | import sys | |
| 7 | import tempfile | |
| 8 | import unittest | |
| 9 | from unittest import mock | |
| 10 | ||
| 11 | HERE = os.path.dirname(os.path.abspath(__file__)) | |
| 12 | sys.path.insert(0, os.path.join(HERE, "..", "lib")) | |
| 13 | import usage # noqa: E402 | |
| 14 | ||
| 15 | FIXTURE = os.path.join(HERE, "fixtures", "usage", "sess1.jsonl") | |
| 16 | ||
| 17 | ||
| 18 | class Trailers(unittest.TestCase): | |
| 19 | def lines(self): | |
| 20 | with mock.patch.object(usage.subprocess, "run", side_effect=FileNotFoundError): # no rtk: its lines are omitted | |
| 21 | return usage.trailers(FIXTURE) | |
| 22 | ||
| 23 | def test_exact_output(self): | |
| 24 | # Duplicate lines per message id count once; <synthetic> is skipped; the subagent file is included; | |
| 25 | # a cache write without the 5m/1h split counts as 5m; models sorted by name. | |
| 26 | self.assertEqual(self.lines(), [ | |
| 27 | "FW-Tokens-In: 76", | |
| 28 | "FW-Tokens-Out: 213", | |
| 29 | "FW-Tokens-Cached: 100970", | |
| 30 | "FW-Tokens-Method: 2", | |
| 31 | "FW-Model-Tokens: claude-haiku-5-5 in=7 out=3 read=70 write5m=0 write1h=0", | |
| 32 | "FW-Model-Tokens: claude-haiku-5-5@over100k in=50 out=20 read=99000 write5m=2000 write1h=0", | |
| 33 | "FW-Model-Tokens: claude-opus-5-5 in=13 out=130 read=1300 write5m=0 write1h=2050", | |
| 34 | "FW-Model-Tokens: claude-sonnet-5-5 in=6 out=60 read=600 write5m=500 write1h=400", | |
| 35 | "FW-Session: sess1", | |
| 36 | ]) | |
| 37 | ||
| 38 | def test_over_100k_prompt_is_filed_separately(self): | |
| 39 | _, _, _, _, per = usage.scan(FIXTURE) | |
| 40 | self.assertEqual(per["claude-haiku-5-5@over100k"], [50, 20, 99000, 2000, 0]) | |
| 41 | self.assertEqual(per["claude-haiku-5-5"], [7, 3, 70, 0, 0]) | |
| 42 | self.assertNotIn("<synthetic>", [k for k, v in per.items() if any(v)]) | |
| 43 | ||
| 44 | def test_rtk_lines_when_rtk_works(self): | |
| 45 | out = '{"summary": {"total_saved": 12, "total_commands": 3}}' | |
| 46 | done = subprocess.CompletedProcess([], 0, stdout=out, stderr="") | |
| 47 | with mock.patch.object(usage.subprocess, "run", return_value=done): | |
| 48 | lines = usage.trailers(FIXTURE) | |
| 49 | self.assertEqual(lines[-2:], ["FW-Rtk-Saved: 12", "FW-Rtk-Cmds: 3"]) | |
| 50 | ||
| 51 | def test_cli_names_session_and_warns_about_old_file(self): | |
| 52 | with tempfile.TemporaryDirectory() as root: | |
| 53 | proj = os.path.join(root, "-p") # the project dir name for cwd /p | |
| 54 | os.makedirs(proj) | |
| 55 | shutil.copy(FIXTURE, os.path.join(proj, "sess1.jsonl")) | |
| 56 | os.utime(os.path.join(proj, "sess1.jsonl"), (0, 0)) | |
| 57 | out, err = io.StringIO(), io.StringIO() | |
| 58 | with mock.patch.object(usage, "root", root), mock.patch.object(usage.os, "getcwd", return_value="/p"), \ | |
| 59 | mock.patch.object(usage, "trailers", return_value=["FW-Tokens-Method: 2"]), \ | |
| 60 | mock.patch("sys.stdout", out), mock.patch("sys.stderr", err): | |
| 61 | usage.main(["--trailers", "--session", "sess1"]) | |
| 62 | self.assertEqual(out.getvalue(), "FW-Tokens-Method: 2\n") | |
| 63 | self.assertIn("sess1.jsonl", err.getvalue()) | |
| 64 | ||
| 65 | class DefaultAndNow(unittest.TestCase): | |
| 66 | """The main() refactor must not change the older outputs.""" | |
| 67 | def run_main(self, args): | |
| 68 | with tempfile.TemporaryDirectory() as root: | |
| 69 | proj = os.path.join(root, "-p") | |
| 70 | os.makedirs(proj) | |
| 71 | shutil.copy(FIXTURE, os.path.join(proj, "sess1.jsonl")) | |
| 72 | shutil.copytree(os.path.join(os.path.dirname(FIXTURE), "sess1"), os.path.join(proj, "sess1")) | |
| 73 | out = io.StringIO() | |
| 74 | with mock.patch.object(usage, "root", root), mock.patch.object(usage.os, "getcwd", return_value="/p"), \ | |
| 75 | mock.patch("sys.stdout", out): | |
| 76 | try: | |
| 77 | usage.main(args) | |
| 78 | except SystemExit: | |
| 79 | pass | |
| 80 | return out.getvalue() | |
| 81 | ||
| 82 | def test_now_line(self): | |
| 83 | self.assertEqual(self.run_main(["--now"]), | |
| 84 | "sess1.js: 6 msgs, context now 0 tokens, cache-read so far 100,970, output 213; " | |
| 85 | "models: haiku-5-5 2, opus-5-5 1, sonnet-5-5 2\n") | |
| 86 | ||
| 87 | def test_default_table(self): | |
| 88 | lines = self.run_main([]).splitlines() | |
| 89 | self.assertTrue(lines[0].startswith("session project")) | |
| 90 | self.assertRegex(lines[1], r"^sess1\.js +-p\s+6\s+0\s+100,970\s+213 \d{4}-\d\d-\d\d \d\d:\d\d$") | |
| 91 | ||
| 92 | def test_session_without_id_is_a_usage_message(self): | |
| 93 | out = self.run_main(["--trailers", "--session"]) | |
| 94 | self.assertEqual(out, "") # exits with the usage message, no traceback | |
| 95 | with mock.patch.object(usage.sys, "exit", side_effect=SystemExit("usage: x")), \ | |
| 96 | mock.patch.object(usage.glob, "glob", return_value=[FIXTURE]): | |
| 97 | with self.assertRaises(SystemExit) as e: | |
| 98 | usage.main(["--trailers", "--session"]) | |
| 99 | self.assertIn("usage", str(e.exception)) | |
| 100 | ||
| 101 | ||
| 102 | if __name__ == "__main__": | |
| 103 | unittest.main() |