107 lines
4.1 KiB
Python
107 lines
4.1 KiB
Python
"""Tests for count_tokens.py: output schema, the forced fallback, files, directories and stdin."""
|
|
|
|
import builtins
|
|
import importlib.util
|
|
import json
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
SCRIPT = Path(__file__).resolve().parents[1] / "count_tokens.py"
|
|
|
|
SAMPLE = (
|
|
"The builder is platform-agnostic. Nothing assumes a single runtime, and no "
|
|
"model list is ever hardcoded. Token counts replace line counts as the one "
|
|
"length metric, with a chars-over-four fallback when tiktoken is absent.\n"
|
|
) * 8
|
|
|
|
|
|
def load_module():
|
|
spec = importlib.util.spec_from_file_location("count_tokens", SCRIPT)
|
|
mod = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(mod)
|
|
return mod
|
|
|
|
|
|
def run(*args, stdin=None):
|
|
return subprocess.run([sys.executable, str(SCRIPT), *args], input=stdin, capture_output=True, text=True)
|
|
|
|
|
|
class CountTokensTest(unittest.TestCase):
|
|
def test_count_returns_int_and_known_method(self):
|
|
tokens, method = load_module().count_tokens(SAMPLE)
|
|
self.assertIn(method, ("tiktoken", "fallback"))
|
|
self.assertIsInstance(tokens, int)
|
|
self.assertGreater(tokens, 0)
|
|
if method == "fallback":
|
|
self.assertEqual(tokens, len(SAMPLE) // 4)
|
|
|
|
def test_fallback_when_tiktoken_import_blocked(self):
|
|
mod = load_module()
|
|
real_import = builtins.__import__
|
|
|
|
def blocked(name, *args, **kwargs):
|
|
if name == "tiktoken" or name.startswith("tiktoken."):
|
|
raise ImportError("blocked for test")
|
|
return real_import(name, *args, **kwargs)
|
|
|
|
builtins.__import__ = blocked
|
|
try:
|
|
tokens, method = mod.count_tokens(SAMPLE)
|
|
finally:
|
|
builtins.__import__ = real_import
|
|
self.assertEqual(method, "fallback")
|
|
self.assertEqual(tokens, len(SAMPLE) // 4)
|
|
|
|
def test_fallback_within_tolerance_of_tiktoken(self):
|
|
mod = load_module()
|
|
real_tokens, method = mod.count_tokens(SAMPLE)
|
|
if method != "tiktoken":
|
|
self.skipTest("tiktoken not installed")
|
|
fallback = len(SAMPLE) // 4
|
|
self.assertTrue(real_tokens * 0.5 <= fallback <= real_tokens * 1.5)
|
|
|
|
def test_cli_file_schema(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
f = Path(tmp) / "sample.md"
|
|
f.write_text(SAMPLE, encoding="utf-8")
|
|
res = run(str(f))
|
|
self.assertEqual(res.returncode, 0, res.stderr)
|
|
data = json.loads(res.stdout)
|
|
self.assertEqual(set(data), {"files", "total", "method"})
|
|
self.assertEqual(data["files"][0]["path"], str(f))
|
|
self.assertEqual(data["total"], data["files"][0]["tokens"])
|
|
self.assertGreater(data["total"], 0)
|
|
|
|
def test_cli_directory_recurses_and_skips_hidden(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
root = Path(tmp)
|
|
(root / "SKILL.md").write_text(SAMPLE, encoding="utf-8")
|
|
(root / "references").mkdir()
|
|
(root / "references" / "a.md").write_text(SAMPLE, encoding="utf-8")
|
|
(root / ".memlog.md").write_text(SAMPLE, encoding="utf-8")
|
|
(root / "logo.png").write_bytes(b"\x89PNG")
|
|
data = json.loads(run(str(root)).stdout)
|
|
paths = [Path(f["path"]).name for f in data["files"]]
|
|
self.assertEqual(paths, ["SKILL.md", "a.md"])
|
|
self.assertEqual(data["total"], sum(f["tokens"] for f in data["files"]))
|
|
|
|
def test_cli_stdin_matches_file(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
f = Path(tmp) / "sample.md"
|
|
f.write_text(SAMPLE, encoding="utf-8")
|
|
from_file = json.loads(run(str(f)).stdout)
|
|
from_stdin = json.loads(run("--stdin", stdin=SAMPLE).stdout)
|
|
self.assertEqual(from_stdin["files"][0]["path"], "<stdin>")
|
|
self.assertEqual(from_stdin["total"], from_file["total"])
|
|
|
|
def test_cli_requires_input(self):
|
|
self.assertNotEqual(run().returncode, 0)
|
|
self.assertNotEqual(run("/nonexistent/path/x.md").returncode, 0)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|