forked from phoenix/litellm-mirror
(feat) Allow failed DB connection requests to allow virtual keys with allow_failed_db_requests
(#6605)
* fix use helper for _handle_failed_db_connection_for_get_key_object * track ALLOW_FAILED_DB_REQUESTS on prometheus * fix allow_failed_db_requests check * fix allow_requests_on_db_unavailable * fix allow_requests_on_db_unavailable * docs allow_requests_on_db_unavailable * identify user_id as litellm_proxy_admin_name when DB is failing * test_handle_failed_db_connection * fix test_user_api_key_auth_db_unavailable * update best practices for prod doc * update best practices for prod * fix handle db failure
This commit is contained in:
parent
eb171e6d95
commit
e3519aa5ae
6 changed files with 224 additions and 3 deletions
|
@ -12,6 +12,11 @@ sys.path.insert(
|
|||
0, os.path.abspath("../..")
|
||||
) # 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
|
||||
from litellm.proxy._types import LiteLLM_EndUserTable, LiteLLM_BudgetTable
|
||||
|
@ -60,3 +65,33 @@ async def test_get_end_user_object(customer_spend, customer_budget):
|
|||
customer_spend, customer_budget, str(e)
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@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"
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue