291 lines
9.9 KiB
Python
291 lines
9.9 KiB
Python
# #region Test.TokenBudget [C:3] [TYPE Module] [SEMANTICS test, translate, token, budget, estimation]
|
|
# @BRIEF Tests for _token_budget.py — token estimation, batch sizing, output-aware constraints.
|
|
# @RELATION BINDS_TO -> [_token_budget]
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent / "src"))
|
|
|
|
import pytest
|
|
|
|
from src.plugins.translate._token_budget import (
|
|
_estimate_tokens_for_text,
|
|
_count_rows_that_fit,
|
|
_calculate_output_tokens,
|
|
_compute_max_rows_by_output,
|
|
_apply_output_aware_batch_sizing,
|
|
_build_warning,
|
|
estimate_token_budget,
|
|
)
|
|
|
|
|
|
class TestEstimateTokensForText:
|
|
"""_estimate_tokens_for_text — CJK-aware heuristic token counting."""
|
|
|
|
def test_empty_string(self):
|
|
"""Empty text returns 1."""
|
|
assert _estimate_tokens_for_text("") == 1
|
|
|
|
def test_ascii_only(self):
|
|
"""Pure ASCII text ~1.8 chars/token."""
|
|
result = _estimate_tokens_for_text("hello world")
|
|
assert result >= 1
|
|
assert isinstance(result, int)
|
|
|
|
def test_cjk_only(self):
|
|
"""Pure CJK text ~1.0 chars/token."""
|
|
result = _estimate_tokens_for_text("你好世界")
|
|
assert result >= 1
|
|
|
|
def test_mixed_text(self):
|
|
"""Mixed CJK + ASCII."""
|
|
result = _estimate_tokens_for_text("hello 你好 world 世界")
|
|
assert result >= 1
|
|
|
|
def test_cjk_ranges(self):
|
|
"""CJK includes Unicode ranges. Fullwidth forms."""
|
|
result = _estimate_tokens_for_text("\uff01\uff02")
|
|
assert result >= 1
|
|
|
|
def test_long_text(self):
|
|
"""Long text estimate is proportional."""
|
|
short = _estimate_tokens_for_text("hello")
|
|
long_ = _estimate_tokens_for_text("hello " * 100)
|
|
assert long_ > short
|
|
|
|
|
|
class TestCountRowsThatFit:
|
|
"""_count_rows_that_fit — consecutive rows within budget."""
|
|
|
|
def test_all_rows_fit(self):
|
|
"""All rows fit within budget."""
|
|
safe, total = _count_rows_that_fit([10, 20, 30], 5000)
|
|
assert safe == 3
|
|
assert total == 60
|
|
|
|
def test_partial_fit(self):
|
|
"""Only first 2 rows fit."""
|
|
safe, total = _count_rows_that_fit([10, 20, 5000], 2100)
|
|
assert safe == 2
|
|
assert total == 30
|
|
|
|
def test_no_rows_fit(self):
|
|
"""First row exceeds budget."""
|
|
safe, total = _count_rows_that_fit([5000], 1000)
|
|
assert safe == 0
|
|
assert total == 0
|
|
|
|
def test_empty_input(self):
|
|
"""Empty list returns (0, 0)."""
|
|
safe, total = _count_rows_that_fit([], 5000)
|
|
assert safe == 0
|
|
assert total == 0
|
|
|
|
def test_single_row_just_fits(self):
|
|
"""Single row just fits (with reasoning overhead)."""
|
|
budget = 200 + 2001 # row tokens + REASONING_OVERHEAD + 1
|
|
safe, total = _count_rows_that_fit([200], budget)
|
|
assert safe == 1
|
|
assert total == 200
|
|
|
|
|
|
class TestCalculateOutputTokens:
|
|
"""_calculate_output_tokens — output budget estimation."""
|
|
|
|
def test_basic(self):
|
|
"""Basic calculation returns positive int."""
|
|
result = _calculate_output_tokens(
|
|
safe_size=5, num_languages=2
|
|
)
|
|
assert result >= 1
|
|
assert isinstance(result, int)
|
|
|
|
def test_zero_safe_size(self):
|
|
"""Zero safe_size returns only overhead."""
|
|
result = _calculate_output_tokens(safe_size=0, num_languages=2)
|
|
assert result >= 2000 # REASONING_OVERHEAD
|
|
|
|
def test_large_batch(self):
|
|
"""Large batch produces larger estimate."""
|
|
small = _calculate_output_tokens(1, 1)
|
|
large = _calculate_output_tokens(50, 5)
|
|
assert large > small
|
|
|
|
|
|
class TestComputeMaxRowsByOutput:
|
|
"""_compute_max_rows_by_output — output-constrained row limit."""
|
|
|
|
def test_basic(self):
|
|
"""Positive result for normal params."""
|
|
result = _compute_max_rows_by_output(max_output_tokens=16384, num_languages=2)
|
|
assert result >= 1
|
|
|
|
def test_small_output_budget(self):
|
|
"""Tiny output budget returns 1."""
|
|
result = _compute_max_rows_by_output(max_output_tokens=100, num_languages=2)
|
|
assert result >= 1
|
|
|
|
def test_zero_per_row(self):
|
|
"""Edge: zero per_row returns fallback 20."""
|
|
result = _compute_max_rows_by_output(max_output_tokens=16384, num_languages=0)
|
|
assert result >= 1
|
|
|
|
|
|
class TestApplyOutputAwareBatchSizing:
|
|
"""_apply_output_aware_batch_sizing — reduce batch until output fits."""
|
|
|
|
def test_no_reduction_needed(self):
|
|
"""Batch fits without reduction."""
|
|
result = _apply_output_aware_batch_sizing(
|
|
safe_size=10, num_languages=1, max_output_tokens=50000
|
|
)
|
|
assert result == 10
|
|
|
|
def test_full_reduction(self):
|
|
"""Batch reduced when output exceeds budget."""
|
|
result = _apply_output_aware_batch_sizing(
|
|
safe_size=10, num_languages=3, max_output_tokens=2000
|
|
)
|
|
assert result < 10
|
|
|
|
def test_zero_size(self):
|
|
"""Zero safe_size returns 0."""
|
|
result = _apply_output_aware_batch_sizing(
|
|
safe_size=0, num_languages=1, max_output_tokens=50000
|
|
)
|
|
assert result == 0
|
|
|
|
|
|
class TestBuildWarning:
|
|
"""_build_warning — warning message generation."""
|
|
|
|
def test_no_warning_when_size_ok(self):
|
|
"""No warning when batch size matches safe size."""
|
|
assert _build_warning(10, 10, 20, 64000, 1000, 500, None) is None
|
|
|
|
def test_reduced_batch_warning(self):
|
|
"""Warning when safe_size < requested batch_size."""
|
|
w = _build_warning(10, 5, 20, 64000, 1000, 500, None)
|
|
assert w is not None
|
|
assert "Reduced" in w
|
|
|
|
def test_auto_calc_warning(self):
|
|
"""Warning when auto-calculated batch smaller than total rows."""
|
|
w = _build_warning(None, 5, 20, 64000, 1000, 500, None)
|
|
assert w is not None
|
|
assert "Auto-calculated" in w
|
|
|
|
def test_dict_warning_appended(self):
|
|
"""Dictionary warning is appended."""
|
|
w = _build_warning(10, 5, 20, 64000, 1000, 500, "Dict entries capped")
|
|
assert w is not None
|
|
assert "Dict entries capped" in w
|
|
|
|
def test_dict_warning_only(self):
|
|
"""Dictionary warning alone when no batch reduction."""
|
|
w = _build_warning(None, 20, 20, 64000, 1000, 500, "Dict capped")
|
|
assert w == "Dict capped"
|
|
|
|
|
|
class TestEstimateTokenBudget:
|
|
"""estimate_token_budget — main budget estimation entry point."""
|
|
|
|
def test_empty_source_rows(self):
|
|
"""Empty source rows returns sensible defaults."""
|
|
result = estimate_token_budget(
|
|
source_rows=[],
|
|
target_languages=["ru"],
|
|
)
|
|
assert result["batch_size_adjusted"] >= 1
|
|
assert "estimated_input_tokens" in result
|
|
assert "estimated_output_tokens" in result
|
|
assert "max_output_needed" in result
|
|
|
|
def test_single_row_single_lang(self):
|
|
"""Single row, single language."""
|
|
result = estimate_token_budget(
|
|
source_rows=[{"source_text": "hello world"}],
|
|
target_languages=["ru"],
|
|
)
|
|
assert result["batch_size_adjusted"] >= 1
|
|
|
|
def test_multiple_rows(self):
|
|
"""Multiple rows produce batch estimate."""
|
|
rows = [{"source_text": f"row {i}"} for i in range(5)]
|
|
result = estimate_token_budget(
|
|
source_rows=rows,
|
|
target_languages=["ru", "de"],
|
|
max_output_tokens=32000,
|
|
)
|
|
assert result["batch_size_adjusted"] >= 1
|
|
|
|
def test_with_context_columns(self):
|
|
"""Context columns add to token estimate."""
|
|
rows = [{"source_text": "hello", "description": "a long description here"}]
|
|
result = estimate_token_budget(
|
|
source_rows=rows,
|
|
target_languages=["ru"],
|
|
source_column="source_text",
|
|
context_columns=["description"],
|
|
)
|
|
assert result["batch_size_adjusted"] >= 1
|
|
|
|
def test_with_dictionary_entries(self):
|
|
"""Dictionary entries add tokens."""
|
|
rows = [{"source_text": "hello"}]
|
|
result = estimate_token_budget(
|
|
source_rows=rows,
|
|
target_languages=["ru"],
|
|
dictionary_entries=[{"id": 1}, {"id": 2}],
|
|
)
|
|
assert result["batch_size_adjusted"] >= 1
|
|
|
|
def test_with_provider_info(self):
|
|
"""Provider info resolves context window."""
|
|
rows = [{"source_text": "hello"}]
|
|
result = estimate_token_budget(
|
|
source_rows=rows,
|
|
target_languages=["ru"],
|
|
provider_info="gpt-4o-mini",
|
|
)
|
|
assert result["batch_size_adjusted"] >= 1
|
|
|
|
def test_explicit_context_window(self):
|
|
"""Explicit context_window overrides provider default."""
|
|
rows = [{"source_text": "hello"}]
|
|
result = estimate_token_budget(
|
|
source_rows=rows,
|
|
target_languages=["ru"],
|
|
context_window=128000,
|
|
max_output_tokens=4096,
|
|
)
|
|
assert result["batch_size_adjusted"] >= 1
|
|
assert result["available_input_budget"] == 128000 - 4096
|
|
|
|
def test_large_batch_reduced(self):
|
|
"""Batch size reduced when too large."""
|
|
many_rows = [{"source_text": "x" * 5000} for _ in range(100)]
|
|
result = estimate_token_budget(
|
|
source_rows=many_rows,
|
|
target_languages=["ru", "de", "fr"],
|
|
batch_size=50,
|
|
)
|
|
assert result["batch_size_adjusted"] >= 1
|
|
|
|
def test_no_target_languages_defaults(self):
|
|
"""Empty target_languages defaults to ['en']."""
|
|
result = estimate_token_budget(
|
|
source_rows=[{"source_text": "hello"}],
|
|
target_languages=[],
|
|
)
|
|
assert result["batch_size_adjusted"] >= 1
|
|
|
|
def test_provider_not_found_fallback(self):
|
|
"""Unknown provider falls back to defaults."""
|
|
result = estimate_token_budget(
|
|
source_rows=[{"source_text": "hello"}],
|
|
target_languages=["ru"],
|
|
provider_info="unknown-model-42",
|
|
)
|
|
assert result["batch_size_adjusted"] >= 1
|