Log in
benchmarksPublic

Viewing files at commit 878fc46b2f31Commit detailsView latest

Browse files
tests/test_token_calibration_campaign.py 115 lines · 4.8 KB · python
1
"""Tests for the token calibration campaign CLI (heuristic method only;
2
network methods are never exercised in tests)."""
3
4
from __future__ import annotations
5
6
import json
7
import subprocess
8
import sys
9
import tempfile
10
import unittest
11
from pathlib import Path
12
13
REPO_ROOT = Path(__file__).resolve().parents[1]
14
SCRIPT = REPO_ROOT / "scripts" / "token_calibration_campaign.py"
15
16
sys.path.insert(0, str(REPO_ROOT / "scripts"))
17
18
from oakbench.calibration import calibrated_tokens, load_calibration
19
20
21
def run_campaign(samples_dir: Path, out: Path, *extra: str) -> subprocess.CompletedProcess[str]:
22
    return subprocess.run(
23
        [
24
            sys.executable,
25
            str(SCRIPT),
26
            "--samples-dir",
27
            str(samples_dir),
28
            "--method",
29
            "heuristic",
30
            "--out",
31
            str(out),
32
            "--calibration-version",
33
            "cal-test-v1",
34
            *extra,
35
        ],
36
        capture_output=True,
37
        text=True,
38
    )
39
40
41
class CampaignTest(unittest.TestCase):
42
    def write_samples(self, samples_dir: Path) -> None:
43
        samples_dir.mkdir()
44
        prose = "The quick brown fox jumps over the lazy dog near the river bank. "
45
        diff = "@@ -1,4 +1,4 @@\n-    return cents - percent\n+    return cents - 1\n"
46
        for index in range(1, 4):
47
            (samples_dir / f"status_output_{index:02d}.txt").write_text(prose * (index + 2))
48
            (samples_dir / f"diff_output_{index:02d}.txt").write_text(diff * (index + 2))
49
50
    def test_two_measured_styles_and_two_unmeasured(self) -> None:
51
        with tempfile.TemporaryDirectory() as tmp:
52
            samples_dir = Path(tmp) / "samples"
53
            self.write_samples(samples_dir)
54
            out = Path(tmp) / "token_calibration.json"
55
            result = run_campaign(samples_dir, out)
56
            self.assertEqual(result.returncode, 0, msg=result.stderr)
57
58
            document = json.loads(out.read_text())
59
            self.assertEqual(document["schema_version"], 1)
60
            self.assertEqual(document["calibration_version"], "cal-test-v1")
61
            self.assertEqual(document["method"], "composition_heuristic")
62
            self.assertEqual(document["samples_per_style"], 3)
63
64
            for style, expected_factor in (("status_output", 4.0), ("diff_output", 3.2)):
65
                entry = document["factors"][style]
66
                self.assertEqual(entry["n"], 3)
67
                self.assertAlmostEqual(entry["chars_per_token"], expected_factor, places=6)
68
                self.assertIsNotNone(entry["ci_low"])
69
                self.assertIsNotNone(entry["ci_high"])
70
                self.assertLessEqual(entry["ci_low"], entry["chars_per_token"])
71
                self.assertLessEqual(entry["chars_per_token"], entry["ci_high"])
72
73
            # Styles with no samples keep null factors with a reason (ADR-0002).
74
            for style in ("log_output", "prose"):
75
                entry = document["factors"][style]
76
                self.assertEqual(entry["n"], 0)
77
                self.assertIsNone(entry["chars_per_token"])
78
                self.assertIsNone(entry["ci_low"])
79
                self.assertIsNone(entry["ci_high"])
80
                self.assertTrue(entry["reason"])
81
82
            # Output round-trips through the loader and prices reports.
83
            calibration = load_calibration(out)
84
            priced = calibrated_tokens(400, "status_output", calibration)
85
            self.assertAlmostEqual(priced["tokens_estimated"], 100.0, places=6)
86
            self.assertEqual(priced["calibration_version"], "cal-test-v1")
87
            self.assertIsNone(calibrated_tokens(400, "prose", calibration))
88
89
    def test_trailing_newline_and_sorted_keys(self) -> None:
90
        with tempfile.TemporaryDirectory() as tmp:
91
            samples_dir = Path(tmp) / "samples"
92
            self.write_samples(samples_dir)
93
            out = Path(tmp) / "token_calibration.json"
94
            result = run_campaign(samples_dir, out)
95
            self.assertEqual(result.returncode, 0, msg=result.stderr)
96
            text = out.read_text()
97
            self.assertTrue(text.endswith("}\n"))
98
            self.assertEqual(text, json.dumps(json.loads(text), indent=2, sort_keys=True) + "\n")
99
100
    def test_no_styles_measured_exits_one(self) -> None:
101
        with tempfile.TemporaryDirectory() as tmp:
102
            samples_dir = Path(tmp) / "samples"
103
            samples_dir.mkdir()
104
            out = Path(tmp) / "token_calibration.json"
105
            result = run_campaign(samples_dir, out)
106
            self.assertEqual(result.returncode, 1)
107
            document = json.loads(out.read_text())
108
            self.assertEqual(document["samples_per_style"], 0)
109
            for entry in document["factors"].values():
110
                self.assertIsNone(entry["chars_per_token"])
111
                self.assertEqual(entry["n"], 0)
112
113
114
if __name__ == "__main__":
115
    unittest.main()