wehub-resource-sync
150 lines
5.8 KiB
Python
150 lines
5.8 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Tests for Anthropic context scaling functionality."""
|
|
|
|
from dataclasses import dataclass, field
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from omlx.settings import ClaudeCodeSettings, GlobalSettings
|
|
|
|
|
|
@dataclass
|
|
class MockServerState:
|
|
"""Mock server state for testing."""
|
|
|
|
global_settings: GlobalSettings | None = None
|
|
sampling: MagicMock = field(default_factory=lambda: MagicMock(max_context_window=32768))
|
|
settings_manager: object | None = None
|
|
|
|
|
|
class TestScaleAnthropicTokens:
|
|
"""Tests for scale_anthropic_tokens function."""
|
|
|
|
def _make_server_state(
|
|
self,
|
|
enabled: bool = False,
|
|
target: int = 200000,
|
|
max_context: int = 32768,
|
|
) -> MockServerState:
|
|
"""Create a mock server state with given settings."""
|
|
gs = GlobalSettings.__new__(GlobalSettings)
|
|
gs.claude_code = ClaudeCodeSettings(
|
|
context_scaling_enabled=enabled,
|
|
target_context_size=target,
|
|
)
|
|
state = MockServerState(global_settings=gs)
|
|
state.sampling = MagicMock(max_context_window=max_context)
|
|
state.settings_manager = None
|
|
return state
|
|
|
|
def test_scaling_disabled_returns_original(self):
|
|
"""When scaling is disabled, return original token count."""
|
|
state = self._make_server_state(enabled=False, target=200000, max_context=32768)
|
|
with (
|
|
patch("omlx.server._server_state", state),
|
|
patch("omlx.server.get_max_context_window", return_value=32768),
|
|
):
|
|
from omlx.server import scale_anthropic_tokens
|
|
|
|
result = scale_anthropic_tokens(1000)
|
|
assert result == 1000
|
|
|
|
def test_scaling_enabled_basic(self):
|
|
"""Test basic scaling: 32k model, 200k target."""
|
|
state = self._make_server_state(enabled=True, target=200000, max_context=100000)
|
|
with (
|
|
patch("omlx.server._server_state", state),
|
|
patch("omlx.server.get_max_context_window", return_value=100000),
|
|
):
|
|
from omlx.server import scale_anthropic_tokens
|
|
|
|
# 50000 * (200000 / 100000) = 100000
|
|
result = scale_anthropic_tokens(50000, "test-model")
|
|
assert result == 100000
|
|
|
|
def test_scaling_preserves_ratio(self):
|
|
"""Test that scaling preserves usage ratio."""
|
|
state = self._make_server_state(enabled=True, target=200000, max_context=131072)
|
|
with (
|
|
patch("omlx.server._server_state", state),
|
|
patch("omlx.server.get_max_context_window", return_value=131072),
|
|
):
|
|
from omlx.server import scale_anthropic_tokens
|
|
|
|
# 96k used on 128k model = 75%
|
|
# 96k used on 128k model = 73.2%
|
|
# scaled: 96000 * (200000/131072) = 146484
|
|
result = scale_anthropic_tokens(96000, "test-model")
|
|
expected = int(96000 * 200000 / 131072)
|
|
assert result == expected
|
|
|
|
def test_no_scaling_when_actual_ge_target(self):
|
|
"""When actual context >= target, no scaling needed."""
|
|
state = self._make_server_state(enabled=True, target=100000, max_context=200000)
|
|
with (
|
|
patch("omlx.server._server_state", state),
|
|
patch("omlx.server.get_max_context_window", return_value=200000),
|
|
):
|
|
from omlx.server import scale_anthropic_tokens
|
|
|
|
result = scale_anthropic_tokens(50000, "test-model")
|
|
assert result == 50000
|
|
|
|
def test_no_scaling_when_actual_equals_target(self):
|
|
"""When actual context == target, no scaling."""
|
|
state = self._make_server_state(enabled=True, target=200000, max_context=250000)
|
|
with (
|
|
patch("omlx.server._server_state", state),
|
|
patch("omlx.server.get_max_context_window", return_value=250000),
|
|
):
|
|
from omlx.server import scale_anthropic_tokens
|
|
|
|
result = scale_anthropic_tokens(100000, "test-model")
|
|
assert result == 100000
|
|
|
|
def test_global_settings_none_returns_original(self):
|
|
"""When global_settings is None, return original."""
|
|
state = MockServerState(global_settings=None)
|
|
with patch("omlx.server._server_state", state):
|
|
from omlx.server import scale_anthropic_tokens
|
|
|
|
result = scale_anthropic_tokens(5000, "test-model")
|
|
assert result == 5000
|
|
|
|
def test_actual_context_none_returns_original(self):
|
|
"""When get_max_context_window returns None, return original."""
|
|
state = self._make_server_state(enabled=True, target=200000, max_context=32768)
|
|
with (
|
|
patch("omlx.server._server_state", state),
|
|
patch("omlx.server.get_max_context_window", return_value=None),
|
|
):
|
|
from omlx.server import scale_anthropic_tokens
|
|
|
|
result = scale_anthropic_tokens(5000, "test-model")
|
|
assert result == 5000
|
|
|
|
def test_zero_tokens(self):
|
|
"""Test with zero token count."""
|
|
state = self._make_server_state(enabled=True, target=200000, max_context=32768)
|
|
with (
|
|
patch("omlx.server._server_state", state),
|
|
patch("omlx.server.get_max_context_window", return_value=32768),
|
|
):
|
|
from omlx.server import scale_anthropic_tokens
|
|
|
|
result = scale_anthropic_tokens(0, "test-model")
|
|
assert result == 0
|
|
|
|
def test_scaling_returns_int(self):
|
|
"""Test that scaling always returns integer."""
|
|
state = self._make_server_state(enabled=True, target=200000, max_context=131072)
|
|
with (
|
|
patch("omlx.server._server_state", state),
|
|
patch("omlx.server.get_max_context_window", return_value=131072),
|
|
):
|
|
from omlx.server import scale_anthropic_tokens
|
|
|
|
result = scale_anthropic_tokens(12345, "test-model")
|
|
assert isinstance(result, int)
|