fuzzyworld-tools / files / tests / test_usage.py
test_usage.py
| 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() |