free-claude-code/tests/core/test_token_estimation.py at main · aqezawat-collab/free-claude-code · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
"""Tests for process-wide plain-text token estimation."""

from unittest.mock import MagicMock, patch

import pytest

from free_claude_code.core import token_estimation


class _RecordingEncoder:
    def __init__(self, tokens: list[int]) -> None:
        self.tokens = tokens
        self.calls: list[tuple[str, tuple[str, ...]]] = []

    def encode(self, text: str, *, disallowed_special: tuple[str, ...]) -> list[int]:
        self.calls.append((text, disallowed_special))
        return self.tokens


def test_encoder_is_loaded_by_canonical_name() -> None:
    encoder = MagicMock()

    with patch.object(
        token_estimation.tiktoken,
        "get_encoding",
        return_value=encoder,
    ) as get_encoding:
        assert token_estimation._load_encoder() is encoder

    get_encoding.assert_called_once_with("cl100k_base")


def test_encoder_acquisition_failure_uses_safe_fallback_warning() -> None:
    with (
        patch.object(
            token_estimation.tiktoken,
            "get_encoding",
            side_effect=RuntimeError("Bearer secret"),
        ),
        patch.object(token_estimation.logger, "warning") as warning,
    ):
        assert token_estimation._load_encoder() is None

    warning.assert_called_once_with(
        "cl100k_base token encoder unavailable ({}); using approximate token estimates",
        "RuntimeError",
    )
    assert "Bearer secret" not in str(warning.call_args)


def test_estimate_text_tokens_uses_cached_encoder() -> None:
    encoder = _RecordingEncoder([1, 2, 3])

    with patch.object(token_estimation, "_ENCODER", encoder):
        assert token_estimation.estimate_text_tokens("hello") == 3

    assert encoder.calls == [("hello", ())]


def test_empty_text_skips_encoder() -> None:
    encoder = _RecordingEncoder([1])

    with patch.object(token_estimation, "_ENCODER", encoder):
        assert token_estimation.estimate_text_tokens("") == 0

    assert encoder.calls == []


@pytest.mark.parametrize(
    ("text", "expected"),
    [
        ("a", 1),
        ("abcdefg", 1),
        ("abcdefgh", 2),
    ],
)
def test_estimate_text_tokens_falls_back_to_character_ratio(
    text: str,
    expected: int,
) -> None:
    with patch.object(token_estimation, "_ENCODER", None):
        assert token_estimation.estimate_text_tokens(text) == expected


def test_encoder_execution_errors_are_not_hidden() -> None:
    encoder = MagicMock()
    encoder.encode.side_effect = RuntimeError("encode failed")

    with (
        patch.object(token_estimation, "_ENCODER", encoder),
        pytest.raises(RuntimeError, match="encode failed"),
    ):
        token_estimation.estimate_text_tokens("hello")