mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
test(proxy): cover get_current_tokens_used and _extract_total_tokens
Exercise the counter-read helper directly (Redis-first hit, cold-counter DB fallback) and the token extractor (payload total, and 0 when the payload is absent). The enforcement tests mock get_current_tokens_used, so these hit the real code paths and lift patch coverage.
This commit is contained in:
parent
dc088d9813
commit
78a7fd1288
1 changed files with 40 additions and 0 deletions
|
|
@ -98,3 +98,43 @@ async def test_exactly_at_cap_passes():
|
|||
)
|
||||
assert result is None
|
||||
assert get_tokens.await_count == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_current_tokens_used_prefers_the_cross_pod_counter():
|
||||
"""A populated tokens:team counter is authoritative over the DB column."""
|
||||
from litellm.proxy.proxy_server import get_current_tokens_used
|
||||
|
||||
async def _counter(counter_key):
|
||||
return 500_000, None
|
||||
|
||||
with patch("litellm.proxy.proxy_server.read_spend_counter_cache_value", _counter):
|
||||
used = await get_current_tokens_used(counter_key="tokens:team:t1", fallback_tokens=7)
|
||||
assert used == 500_000
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_current_tokens_used_falls_back_to_db_when_counter_cold():
|
||||
"""A cold counter (None) falls back to the tokens_used column loaded in auth."""
|
||||
from litellm.proxy.proxy_server import get_current_tokens_used
|
||||
|
||||
async def _cold(counter_key):
|
||||
return None, None
|
||||
|
||||
with patch("litellm.proxy.proxy_server.read_spend_counter_cache_value", _cold):
|
||||
used = await get_current_tokens_used(counter_key="tokens:team:t1", fallback_tokens=1_234)
|
||||
assert used == 1_234
|
||||
|
||||
|
||||
def test_extract_total_tokens_reads_the_standard_logging_payload():
|
||||
"""The counter delta is the StandardLoggingPayload total, the DB write's source."""
|
||||
from litellm.proxy.hooks.proxy_track_cost_callback import _extract_total_tokens
|
||||
|
||||
assert _extract_total_tokens({"total_tokens": 4242}) == 4242
|
||||
|
||||
|
||||
def test_extract_total_tokens_is_zero_without_a_payload():
|
||||
"""No payload leaves the counter untouched rather than raising."""
|
||||
from litellm.proxy.hooks.proxy_track_cost_callback import _extract_total_tokens
|
||||
|
||||
assert _extract_total_tokens(None) == 0
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue