232 lines
9.1 KiB
Python
232 lines
9.1 KiB
Python
|
"""
|
||
|
|
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
|