mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-08 08:08:37 +00:00
test(proxy): cover DB-floor edge cases for global proxy spend
Add regression tests for: - Redis async_set_max returning None (fallback to DB floor). - Unparseable Redis value being treated as 0.0 and floored to DB. Improves patch coverage for the auth-time floor repair path.
This commit is contained in:
parent
e70e4b9aa0
commit
b1b0fd586c
1 changed files with 65 additions and 0 deletions
|
|
@ -6263,6 +6263,71 @@ async def test_global_proxy_spend_marker_cached_db_value_used_without_db_query()
|
|||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_global_proxy_spend_floor_falls_back_when_set_max_returns_none():
|
||||
"""If the monotonic Redis repair returns None (e.g. connection hiccup),
|
||||
the enforced total must still fall back to the authoritative DB floor."""
|
||||
from litellm.proxy.auth.user_api_key_auth import (
|
||||
_fetch_global_spend_with_event_coordination,
|
||||
)
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
|
||||
cache = UserApiKeyCache()
|
||||
redis_cache = MagicMock()
|
||||
redis_cache.async_get_cache = AsyncMock(return_value="2.0")
|
||||
redis_cache.async_set_max = AsyncMock(return_value=None)
|
||||
cache.redis_cache = redis_cache
|
||||
|
||||
proxy_budget_row = MagicMock()
|
||||
proxy_budget_row.spend = 9.0
|
||||
prisma_client = MagicMock()
|
||||
prisma_client.db.litellm_usertable.find_unique = AsyncMock(
|
||||
return_value=proxy_budget_row
|
||||
)
|
||||
|
||||
result = await _fetch_global_spend_with_event_coordination(
|
||||
cache_key="default_user_id:spend",
|
||||
user_api_key_cache=cache,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
|
||||
assert result == 9.0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_global_proxy_spend_treats_unparseable_redis_as_zero():
|
||||
"""A corrupted/non-numeric Redis value is treated as 0.0 and then floored
|
||||
to the DB total so enforcement remains safe."""
|
||||
from litellm.proxy.auth.user_api_key_auth import (
|
||||
_fetch_global_spend_with_event_coordination,
|
||||
)
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
|
||||
cache = UserApiKeyCache()
|
||||
redis_cache = MagicMock()
|
||||
redis_cache.async_get_cache = AsyncMock(return_value="not-a-number")
|
||||
redis_cache.async_set_max = AsyncMock(side_effect=lambda key, value: value)
|
||||
cache.redis_cache = redis_cache
|
||||
|
||||
proxy_budget_row = MagicMock()
|
||||
proxy_budget_row.spend = 7.0
|
||||
prisma_client = MagicMock()
|
||||
prisma_client.db.litellm_usertable.find_unique = AsyncMock(
|
||||
return_value=proxy_budget_row
|
||||
)
|
||||
|
||||
result = await _fetch_global_spend_with_event_coordination(
|
||||
cache_key="default_user_id:spend",
|
||||
user_api_key_cache=cache,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
|
||||
assert result == 7.0
|
||||
redis_cache.async_set_max.assert_awaited_once_with(
|
||||
key="default_user_id:spend", value=7.0
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_temp_budget_increase_applied_for_cached_key():
|
||||
"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue