76 lines
2.2 KiB
Python
76 lines
2.2 KiB
Python
import pytest
|
||
|
||
from genesis.inference.token import (
|
||
_tiktoken_estimator,
|
||
approximate_token_count,
|
||
make_estimator,
|
||
)
|
||
|
||
|
||
def test_approximate_count_minimum():
|
||
assert approximate_token_count("") >= 1
|
||
assert approximate_token_count("a") == 1
|
||
|
||
|
||
def test_approximate_count_linear():
|
||
# 每 4 字符约 1 token（向上取整）
|
||
assert approximate_token_count("abcd") == 1
|
||
assert approximate_token_count("abcdefgh") == 2
|
||
assert approximate_token_count("abcdefghi") == 3
|
||
|
||
|
||
def test_tiktoken_estimator_missing_falls_back():
|
||
# 未安装 tiktoken 或不可用时返回 None（由 make_estimator 回落 approximate）
|
||
r = _tiktoken_estimator("hello")
|
||
assert r is None or isinstance(r, int)
|
||
|
||
|
||
def test_make_estimator_approximate_backend():
|
||
est = make_estimator("approximate")
|
||
assert est("abcd") == 1
|
||
|
||
|
||
def test_make_estimator_default_without_tiktoken(monkeypatch):
|
||
# 强制模拟 tiktoken 缺失：make_estimator 必须回落 approximate
|
||
import builtins
|
||
real_import = builtins.__import__
|
||
|
||
def fake_import(name, *args, **kwargs):
|
||
if name == "tiktoken":
|
||
raise ImportError("no tiktoken")
|
||
return real_import(name, *args, **kwargs)
|
||
|
||
monkeypatch.setattr(builtins, "__import__", fake_import)
|
||
est = make_estimator("tiktoken")
|
||
assert est("abcd") == 1
|
||
|
||
|
||
def test_tiktoken_estimator_success_path(monkeypatch):
|
||
# 注入假 tiktoken 模块：验证编码成功路径（len(encode(text))）
|
||
import sys
|
||
|
||
class _FakeEncoding:
|
||
def encode(self, s: str):
|
||
return ["t"] * len(s)
|
||
|
||
class _FakeTiktoken:
|
||
@staticmethod
|
||
def get_encoding(name):
|
||
assert name == "cl100k_base"
|
||
return _FakeEncoding()
|
||
|
||
monkeypatch.setitem(sys.modules, "tiktoken", _FakeTiktoken)
|
||
assert _tiktoken_estimator("hello") == 5
|
||
|
||
|
||
def test_tiktoken_estimator_encoding_error(monkeypatch):
|
||
# tiktoken 可用但编码抛异常：返回 None（由 make_estimator 回落 approximate）
|
||
import sys
|
||
|
||
class _FakeTiktoken:
|
||
@staticmethod
|
||
def get_encoding(name):
|
||
raise RuntimeError("boom")
|
||
|
||
monkeypatch.setitem(sys.modules, "tiktoken", _FakeTiktoken)
|
||
assert _tiktoken_estimator("hello") is None |