mirror of
https://github.com/Routstr/routstr-core.git
synced 2026-10-05 20:28:23 +00:00
redshift
bc8172783d
_fold_cache_into_input_tokens rolled cache_read/cache_creation tokens into the visible prompt_tokens for every dialect. In the OpenAI family (Venice, OpenAI, DeepSeek, OpenRouter, litellm) prompt_tokens already includes the cached portion, so the client-visible prompt count double-counted cache hits (Venice: 14075-token prompt shown as 27997 after a 13922-token cache read). Billing was unaffected — normalize_usage already subtracts cache exactly once and the fold runs after cost calculation. The fold now mirrors normalize_usage: only Anthropic-native input_tokens (which excludes cache) gets the roll-up; prompt_tokens is left untouched. Adds dict-based regression tests for the Venice/OpenAI, Anthropic-native and litellm-mirror shapes (existing Mock tests never exercised the logic).
266 lines
9.4 KiB
Python
266 lines
9.4 KiB
Python
"""Additional coverage tests for base.py (41% → target 50%+).
|
|
|
|
Tests error message extraction, static helpers, model cache, and cost hooks.
|
|
|
|
These test existing correct behavior — all should PASS.
|
|
"""
|
|
|
|
import json
|
|
from unittest.mock import Mock
|
|
|
|
import pytest
|
|
|
|
from routstr.upstream.base import BaseUpstreamProvider
|
|
|
|
# ===========================================================================
|
|
# _extract_upstream_error_message
|
|
# ===========================================================================
|
|
|
|
def test_extract_error_from_json_body() -> None:
|
|
"""Error message is extracted from JSON upstream error response."""
|
|
p = BaseUpstreamProvider("https://api.test.com", "sk-test")
|
|
body = json.dumps({"error": {"message": "Model not found", "type": "not_found"}}).encode()
|
|
|
|
msg, error_type = p._extract_upstream_error_message(body)
|
|
|
|
assert "Model not found" in msg
|
|
assert error_type == "not_found"
|
|
|
|
|
|
def test_extract_error_from_simple_json() -> None:
|
|
"""Simple JSON error with direct message key."""
|
|
p = BaseUpstreamProvider("https://api.test.com", "sk-test")
|
|
body = json.dumps({"message": "Rate limit exceeded"}).encode()
|
|
|
|
msg, error_type = p._extract_upstream_error_message(body)
|
|
|
|
assert "Rate limit" in msg
|
|
|
|
|
|
def test_extract_error_from_text_body() -> None:
|
|
"""Non-JSON text body is returned as-is."""
|
|
p = BaseUpstreamProvider("https://api.test.com", "sk-test")
|
|
|
|
msg, error_type = p._extract_upstream_error_message(b"Internal Server Error")
|
|
|
|
assert "Internal Server Error" in msg
|
|
|
|
|
|
def test_extract_error_empty_body() -> None:
|
|
"""Empty body returns a generic message."""
|
|
p = BaseUpstreamProvider("https://api.test.com", "sk-test")
|
|
|
|
msg, error_type = p._extract_upstream_error_message(b"")
|
|
|
|
assert isinstance(msg, str)
|
|
assert len(msg) > 0
|
|
|
|
|
|
def test_extract_error_simple_error_string_not_parsed() -> None:
|
|
"""JSON error as plain string (not dict) falls through to generic message."""
|
|
p = BaseUpstreamProvider("https://api.test.com", "sk-test")
|
|
body = json.dumps({"error": "Invalid API key"}).encode()
|
|
|
|
msg, error_type = p._extract_upstream_error_message(body)
|
|
|
|
# Simple error strings not nested in a dict object use generic message
|
|
assert "Upstream request failed" in msg or "Invalid" in msg
|
|
|
|
|
|
# ===========================================================================
|
|
# on_upstream_error_redirect
|
|
# ===========================================================================
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_on_upstream_error_redirect_noop() -> None:
|
|
"""Default implementation is a no-op for non-redirect statuses."""
|
|
p = BaseUpstreamProvider("https://api.test.com", "sk-test")
|
|
await p.on_upstream_error_redirect(402, "Insufficient balance")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_on_upstream_error_redirect_429() -> None:
|
|
"""429 rate limit passes through (subclasses may override)."""
|
|
p = BaseUpstreamProvider("https://api.test.com", "sk-test")
|
|
await p.on_upstream_error_redirect(429, "Rate limited")
|
|
|
|
|
|
# ===========================================================================
|
|
# _fold_cache_into_input_tokens (static method)
|
|
# ===========================================================================
|
|
|
|
def test_fold_cache_no_cache_data() -> None:
|
|
"""Usage without cache details is unchanged."""
|
|
from routstr.upstream.base import BaseUpstreamProvider
|
|
|
|
usage = Mock()
|
|
usage.prompt_tokens = 100
|
|
del usage.prompt_tokens_details # No cache details
|
|
|
|
BaseUpstreamProvider._fold_cache_into_input_tokens(usage)
|
|
# Should not modify the usage object when no cache exists
|
|
|
|
|
|
def test_fold_cache_preserves_total() -> None:
|
|
"""Total prompt tokens remain the same after folding cache."""
|
|
from routstr.upstream.base import BaseUpstreamProvider
|
|
|
|
usage = Mock()
|
|
usage.prompt_tokens = 100
|
|
details = Mock()
|
|
details.cached_tokens = 30
|
|
usage.prompt_tokens_details = details
|
|
|
|
BaseUpstreamProvider._fold_cache_into_input_tokens(usage)
|
|
# prompt_tokens should still be 100 (total unchanged)
|
|
assert usage.prompt_tokens == 100
|
|
|
|
|
|
def test_fold_cache_openai_dialect_prompt_tokens_untouched() -> None:
|
|
"""Venice/OpenAI shape: prompt_tokens already includes cached tokens.
|
|
|
|
Regression: folding cache_read into prompt_tokens double-counted the
|
|
cache (14075 real prompt shown as 27997 after a 13922-token cache read).
|
|
"""
|
|
from routstr.upstream.base import BaseUpstreamProvider
|
|
|
|
usage = {
|
|
"prompt_tokens": 14075,
|
|
"completion_tokens": 24,
|
|
"total_tokens": 14099,
|
|
"prompt_tokens_details": {"cached_tokens": 13922},
|
|
"cache_read_input_tokens": 13922,
|
|
}
|
|
BaseUpstreamProvider._fold_cache_into_input_tokens(usage)
|
|
assert usage["prompt_tokens"] == 14075
|
|
assert usage["cache_read_input_tokens"] == 13922
|
|
|
|
|
|
def test_fold_cache_anthropic_dialect_folds_input_tokens() -> None:
|
|
"""Anthropic-native shape: input_tokens excludes cache, so it is folded."""
|
|
from routstr.upstream.base import BaseUpstreamProvider
|
|
|
|
usage = {
|
|
"input_tokens": 153,
|
|
"output_tokens": 24,
|
|
"cache_read_input_tokens": 13922,
|
|
"cache_creation_input_tokens": 0,
|
|
}
|
|
BaseUpstreamProvider._fold_cache_into_input_tokens(usage)
|
|
assert usage["input_tokens"] == 153 + 13922
|
|
|
|
|
|
def test_fold_cache_litellm_mirror_folds_only_input_tokens() -> None:
|
|
"""Both fields present (litellm mirror): fold input_tokens only."""
|
|
from routstr.upstream.base import BaseUpstreamProvider
|
|
|
|
usage = {
|
|
"prompt_tokens": 14075, # inclusive grand total
|
|
"input_tokens": 153, # additive Anthropic mirror
|
|
"cache_read_input_tokens": 13922,
|
|
}
|
|
BaseUpstreamProvider._fold_cache_into_input_tokens(usage)
|
|
assert usage["prompt_tokens"] == 14075
|
|
assert usage["input_tokens"] == 153 + 13922
|
|
|
|
|
|
# ===========================================================================
|
|
# get_cached_models / get_cached_model_by_id
|
|
# ===========================================================================
|
|
|
|
def test_get_cached_models_returns_list() -> None:
|
|
"""get_cached_models always returns a list."""
|
|
p = BaseUpstreamProvider("https://api.test.com", "sk-test")
|
|
models = p.get_cached_models()
|
|
assert isinstance(models, list)
|
|
|
|
|
|
def test_get_cached_model_by_id_unknown_returns_none() -> None:
|
|
"""Unknown model ID returns None."""
|
|
p = BaseUpstreamProvider("https://api.test.com", "sk-test")
|
|
result = p.get_cached_model_by_id("nonexistent-model-xyz-12345")
|
|
assert result is None
|
|
|
|
|
|
# ===========================================================================
|
|
# get_x_cashu_cost
|
|
# ===========================================================================
|
|
|
|
def test_get_x_cashu_cost_with_usage() -> None:
|
|
"""Cost is calculated from response data with usage info."""
|
|
p = BaseUpstreamProvider("https://api.test.com", "sk-test")
|
|
response_data = {
|
|
"model": "gpt-4",
|
|
"usage": {"prompt_tokens": 100, "completion_tokens": 50},
|
|
}
|
|
|
|
result = p.get_x_cashu_cost(response_data, 100000, None)
|
|
|
|
# Either returns None (needs more data) or a cost object
|
|
assert result is not None
|
|
|
|
|
|
def test_get_x_cashu_cost_no_usage() -> None:
|
|
"""Response without usage returns MaxCostData."""
|
|
p = BaseUpstreamProvider("https://api.test.com", "sk-test")
|
|
response_data = {"model": "gpt-4"}
|
|
|
|
result = p.get_x_cashu_cost(response_data, 100000, None)
|
|
|
|
# Without usage, uses max_cost
|
|
assert result is not None
|
|
|
|
|
|
# ===========================================================================
|
|
# get_balance
|
|
# ===========================================================================
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_balance_raises_not_implemented() -> None:
|
|
"""Default get_balance raises NotImplementedError (no account support)."""
|
|
p = BaseUpstreamProvider("https://api.test.com", "sk-test")
|
|
with pytest.raises(NotImplementedError):
|
|
await p.get_balance()
|
|
|
|
|
|
# ===========================================================================
|
|
# refresh_models_cache
|
|
# ===========================================================================
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_refresh_models_cache_no_providers() -> None:
|
|
"""refresh_models_cache handles empty provider list gracefully."""
|
|
p = BaseUpstreamProvider("https://api.test.com", "sk-test")
|
|
# Default implementation may be a no-op or raise
|
|
try:
|
|
await p.refresh_models_cache()
|
|
except Exception:
|
|
pass # May fail without DB — that's fine
|
|
|
|
|
|
# ===========================================================================
|
|
# fetch_models
|
|
# ===========================================================================
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_fetch_models_returns_list() -> None:
|
|
"""fetch_models returns a model list (or empty) for default provider."""
|
|
p = BaseUpstreamProvider("https://api.test.com", "sk-test")
|
|
try:
|
|
result = await p.fetch_models()
|
|
assert isinstance(result, list)
|
|
except Exception:
|
|
pass # May fail without network
|
|
|
|
|
|
# ===========================================================================
|
|
# create_account
|
|
# ===========================================================================
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_account_raises_not_implemented() -> None:
|
|
"""Default create_account raises NotImplementedError."""
|
|
p = BaseUpstreamProvider("https://api.test.com", "sk-test")
|
|
with pytest.raises(NotImplementedError):
|
|
await p.create_account()
|