| 1 | import tempfile |
| 2 | import unittest |
| 3 | from pathlib import Path |
| 4 | |
| 5 | import scripts.preflight_cost_check as preflight |
| 6 | |
| 7 | |
| 8 | class PreflightCostCheckTests(unittest.TestCase): |
| 9 | def _make_file(self, base: Path, name: str, size: int) -> Path: |
| 10 | p = base / name |
| 11 | p.write_text("x" * size, encoding="utf-8") |
| 12 | return p |
| 13 | |
| 14 | def test_passes_under_cap(self) -> None: |
| 15 | tests_root = Path(__file__).resolve().parent |
| 16 | with tempfile.TemporaryDirectory(dir=tests_root) as tmpdir: |
| 17 | base = Path(tmpdir) |
| 18 | f1 = self._make_file(base, "small.json", 400) |
| 19 | rc = preflight.main(["--context-files", str(f1)]) |
| 20 | self.assertEqual(rc, 0) |
| 21 | |
| 22 | def test_fails_over_cap(self) -> None: |
| 23 | tests_root = Path(__file__).resolve().parent |
| 24 | with tempfile.TemporaryDirectory(dir=tests_root) as tmpdir: |
| 25 | base = Path(tmpdir) |
| 26 | # 2M chars → 500k tokens; at $3/M input that's $1.50 — over cap |
| 27 | f1 = self._make_file(base, "huge.json", 2_000_000) |
| 28 | rc = preflight.main(["--context-files", str(f1)]) |
| 29 | self.assertEqual(rc, 1) |
| 30 | |
| 31 | def test_custom_cap(self) -> None: |
| 32 | tests_root = Path(__file__).resolve().parent |
| 33 | with tempfile.TemporaryDirectory(dir=tests_root) as tmpdir: |
| 34 | base = Path(tmpdir) |
| 35 | f1 = self._make_file(base, "medium.json", 40_000) |
| 36 | # 40k chars → 10k tokens; at $3/M input = $0.03; cap $0.01 should fail |
| 37 | rc = preflight.main(["--context-files", str(f1), "--hard-cap", "0.01"]) |
| 38 | self.assertEqual(rc, 1) |
| 39 | |
| 40 | def test_unknown_model_fails(self) -> None: |
| 41 | tests_root = Path(__file__).resolve().parent |
| 42 | with tempfile.TemporaryDirectory(dir=tests_root) as tmpdir: |
| 43 | base = Path(tmpdir) |
| 44 | f1 = self._make_file(base, "a.json", 100) |
| 45 | rc = preflight.main(["--context-files", str(f1), "--model", "unknown-model"]) |
| 46 | self.assertEqual(rc, 1) |
| 47 | |
| 48 | def test_multiple_files_summed(self) -> None: |
| 49 | tests_root = Path(__file__).resolve().parent |
| 50 | with tempfile.TemporaryDirectory(dir=tests_root) as tmpdir: |
| 51 | base = Path(tmpdir) |
| 52 | f1 = self._make_file(base, "a.json", 200) |
| 53 | f2 = self._make_file(base, "b.json", 200) |
| 54 | tokens = preflight.estimate_input_tokens([f1, f2]) |
| 55 | self.assertEqual(tokens, 100) # 400 chars / 4 = 100 tokens |
| 56 | |
| 57 | def test_missing_file_counts_zero(self) -> None: |
| 58 | tokens = preflight.estimate_input_tokens([Path("/nonexistent/path.json")]) |
| 59 | self.assertEqual(tokens, 0) |
| 60 | |
| 61 | |
| 62 | if __name__ == "__main__": |
| 63 | unittest.main() |