phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-25 02:34:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

handle failed db connections

Browse source
This commit is contained in:
Ishaan Jaff 2025-03-25 23:14:44 -07:00
parent 0a401ee468
commit 6e5d2b1ac7
1 changed files with 0 additions and 33 deletions
Download patch file Download diff file Expand all files Collapse all files

33
tests/proxy_unit_tests/test_auth_checks.py
View file

@ -13,9 +13,6 @@ sys.path.insert(
) # Adds the parent directory to the system path
import pytest, litellm
import httpx
from litellm.proxy.auth.auth_checks import (
_handle_failed_db_connection_for_get_key_object,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.auth_checks import get_end_user_object
from litellm.caching.caching import DualCache
@ -78,36 +75,6 @@ async def test_get_end_user_object(customer_spend, customer_budget):
)
@pytest.mark.asyncio
async def test_handle_failed_db_connection():
"""
Test cases:
1. When allow_requests_on_db_unavailable=True -> return UserAPIKeyAuth
2. When allow_requests_on_db_unavailable=False -> raise original error
"""
from litellm.proxy.proxy_server import general_settings, litellm_proxy_admin_name
# Test case 1: allow_requests_on_db_unavailable=True
general_settings["allow_requests_on_db_unavailable"] = True
mock_error = httpx.ConnectError("Failed to connect to DB")
result = await _handle_failed_db_connection_for_get_key_object(e=mock_error)
assert isinstance(result, UserAPIKeyAuth)
assert result.key_name == "failed-to-connect-to-db"
assert result.token == "failed-to-connect-to-db"
assert result.user_id == litellm_proxy_admin_name
# Test case 2: allow_requests_on_db_unavailable=False
general_settings["allow_requests_on_db_unavailable"] = False
with pytest.raises(httpx.ConnectError) as exc_info:
await _handle_failed_db_connection_for_get_key_object(e=mock_error)
print("_handle_failed_db_connection_for_get_key_object got exception", exc_info)
assert str(exc_info.value) == "Failed to connect to DB"
@pytest.mark.parametrize(
"model, expect_to_work",
[