""" Tests for scripts/token_estimator.py Covers the deterministic token estimator (Phase 2): - HeuristicEstimator chars/4 math incl. empty + boundary lengths - default_estimator() returns a working estimator with no deps / no network - file estimation + graceful decode handling - schema-shaped report wrapper (raw_tokens only in v1) - CLI main() JSON output + exit codes - tiktoken backend (gated by importorskip; relies on vendored cache, never hits the network) sys.path is patched here (no shared conftest) so the test file is self-contained, matching the repo's test style. """ from __future__ import annotations import json import sys from pathlib import Path import pytest # --------------------------------------------------------------------------- # Path bootstrap — inject scripts/ so we can import token_estimator directly. # --------------------------------------------------------------------------- _SCRIPTS_DIR = Path(__file__).parent.parent / "scripts" if str(_SCRIPTS_DIR) not in sys.path: sys.path.insert(0, str(_SCRIPTS_DIR)) from token_estimator import ( # noqa: E402 HeuristicEstimator, TiktokenEstimator, TokenEstimator, default_estimator, main, _O200K_BASE_CACHE_FILENAME, ) # =========================================================================== # HeuristicEstimator — chars/4 math # =========================================================================== class TestHeuristicMath: def test_empty_string_is_zero(self): assert HeuristicEstimator().estimate("") == 0 @pytest.mark.parametrize( "text,expected", [ ("a", 1), # ceil(1/4) ("abc", 1), # ceil(3/4) ("abcd", 1), # boundary: ceil(4/4) == 1 ("abcde", 2), # boundary: ceil(5/4) == 2 ("abcdefgh", 2), # ceil(8/4) == 2 ("a" * 9, 3), # ceil(9/4) == 3 ], ) def test_boundary_lengths(self, text, expected): assert HeuristicEstimator().estimate(text) == expected def test_count_is_nonnegative_and_monotone(self): est = HeuristicEstimator() prev = -1 for n in range(0, 50): count = est.estimate("x" * n) assert count >= 0 assert count >= prev prev = count def test_name_is_heuristic(self): assert HeuristicEstimator().name == "heuristic" # =========================================================================== # default_estimator() — always returns a working estimator, no network # =========================================================================== class TestDefaultEstimator: def test_returns_token_estimator_instance(self, tmp_path): # Empty cache dir => no vendored vocab => must be heuristic, never network. est = default_estimator(cache_dir=tmp_path) assert isinstance(est, TokenEstimator) assert isinstance(est, HeuristicEstimator) assert est.name == "heuristic" def test_works_with_no_deps(self, tmp_path): # The returned estimator must actually estimate, regardless of backend. est = default_estimator(cache_dir=tmp_path) assert est.estimate("hello world") > 0 def test_never_raises_on_missing_cache(self, tmp_path): # Nonexistent cache dir must not raise — falls back cleanly. missing = tmp_path / "does-not-exist" est = default_estimator(cache_dir=missing) assert isinstance(est, HeuristicEstimator) # =========================================================================== # File estimation + decode handling # =========================================================================== class TestFileEstimation: def test_estimate_file_counts_contents(self, tmp_path): f = tmp_path / "doc.md" f.write_text("abcdefgh", encoding="utf-8") # 8 chars -> 2 tokens assert HeuristicEstimator().estimate_file(f) == 2 def test_missing_file_yields_zero(self, tmp_path): assert HeuristicEstimator().estimate_file(tmp_path / "nope.md") == 0 def test_invalid_utf8_decodes_gracefully(self, tmp_path): f = tmp_path / "binary.md" # Invalid UTF-8 bytes — must not raise, replacement chars are counted. f.write_bytes(b"\xff\xfe\xfa\xfbhello") count = HeuristicEstimator().estimate_file(f) assert count > 0 def test_empty_file_yields_zero(self, tmp_path): f = tmp_path / "empty.md" f.write_text("", encoding="utf-8") assert HeuristicEstimator().estimate_file(f) == 0 # =========================================================================== # Schema-shaped report wrapper # =========================================================================== class TestReportShape: def test_estimate_for_report_v1_shape(self): est = HeuristicEstimator() out = est.estimate_for_report("abcd") # 1 token assert out == { "raw_tokens": 1, "injection_frequency": None, "weighted_tokens": None, } def test_raw_tokens_matches_estimate(self): est = HeuristicEstimator() text = "some longer markdown content here" assert est.estimate_for_report(text)["raw_tokens"] == est.estimate(text) # =========================================================================== # CLI main() # =========================================================================== class TestCLI: def test_main_emits_json_for_file(self, tmp_path, capsys): f = tmp_path / "doc.md" f.write_text("abcdefgh", encoding="utf-8") rc = main([str(f), "--cache-dir", str(tmp_path)]) assert rc == 0 out = json.loads(capsys.readouterr().out) assert out["path"] == str(f) assert out["backend"] == "heuristic" assert out["token_estimate"]["raw_tokens"] == 2 assert out["token_estimate"]["injection_frequency"] is None def test_main_missing_file_exits_nonzero(self, tmp_path, capsys): rc = main([str(tmp_path / "missing.md")]) assert rc == 1 err = json.loads(capsys.readouterr().err) assert "error" in err # =========================================================================== # Cache-filename gotcha — guards the offline contract # =========================================================================== class TestCacheFilename: def test_cache_filename_is_sha1_of_blobpath(self): # If this changes, the vendored-vocab pre-check silently misses and # tiktoken would fall to the network. Pin it. assert _O200K_BASE_CACHE_FILENAME == "fb374d419588a4632f3f557e76b4b70aebbca790" # =========================================================================== # tiktoken backend — only if importable; must rely on vendored cache, no net. # =========================================================================== class TestTiktokenBackend: def test_tiktoken_estimator_plausible_counts(self): tiktoken = pytest.importorskip("tiktoken") # Load directly from the real registry (test env, allowed to read its # own installed cache; we assert no-network indirectly by using a # plausible-count check rather than forcing a download). try: encoding = tiktoken.get_encoding("o200k_base") except Exception: pytest.skip("o200k_base vocab unavailable offline in this env") est = TiktokenEstimator(encoding, "o200k_base") assert est.name == "tiktoken:o200k_base" assert est.estimate("") == 0 # A short English sentence is a handful of tokens, fewer than chars. text = "The quick brown fox jumps over the lazy dog." count = est.estimate(text) assert 0 < count < len(text) def test_tiktoken_special_tokens_do_not_raise(self): tiktoken = pytest.importorskip("tiktoken") try: encoding = tiktoken.get_encoding("o200k_base") except Exception: pytest.skip("o200k_base vocab unavailable offline in this env") est = TiktokenEstimator(encoding, "o200k_base") # disallowed_special=() means this special-token-looking text is fine. assert est.estimate("<|endoftext|> in a doc") > 0 def test_default_estimator_selects_tiktoken_when_cache_present(self, tmp_path): tiktoken = pytest.importorskip("tiktoken") # Vendor the vocab into a tmp cache dir under the correct sha1 filename, # then assert default_estimator picks tiktoken (no network: file is local). try: real = tiktoken.get_encoding("o200k_base") # may use installed cache except Exception: pytest.skip("cannot load o200k_base offline in this env") # Best-effort: locate the installed cache file to copy it. src_dir = Path( __import__("os").environ.get("TIKTOKEN_CACHE_DIR", "") ) src = src_dir / _O200K_BASE_CACHE_FILENAME if not src.is_file(): pytest.skip("vendored vocab file not locatable to copy") (tmp_path / _O200K_BASE_CACHE_FILENAME).write_bytes(src.read_bytes()) est = default_estimator(cache_dir=tmp_path) assert est.name == "tiktoken:o200k_base" assert real.encode("hi", disallowed_special=()) is not None