From b560164921b2d993925db8415d96ca806638d100 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Fri, 20 Feb 2026 15:19:35 -0800 Subject: [PATCH] test(auth): cover reconnect-once behavior for key lookup --- .../proxy/auth/test_auth_checks.py | 57 ++++++++++++++++++- 1 file changed, 55 insertions(+), 2 deletions(-) diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 4f8e80c023e..8aad1ebe684 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -10,6 +10,7 @@ sys.path.insert( from datetime import datetime, timedelta +import httpx import pytest import litellm @@ -33,6 +34,7 @@ from litellm.proxy.auth.auth_checks import ( _log_budget_lookup_failure, _virtual_key_max_budget_alert_check, _virtual_key_soft_budget_check, + get_key_object, get_user_object, vector_store_access_check, ) @@ -50,9 +52,10 @@ def set_salt_key(monkeypatch): def reset_constants_module(): """Reset constants module to ensure clean state before each test""" import importlib + from litellm import constants from litellm.proxy.auth import auth_checks - + # Reload modules before test importlib.reload(constants) importlib.reload(auth_checks) @@ -151,6 +154,55 @@ def test_get_key_object_from_ui_hash_key_invalid(): assert key_object is None +@pytest.mark.asyncio +async def test_get_key_object_should_reconnect_once_on_db_connection_error(): + mock_prisma_client = MagicMock() + mock_prisma_client.get_data = AsyncMock( + side_effect=[ + httpx.ConnectError("db connection reset"), + UserAPIKeyAuth(token="hashed-token-1"), + ] + ) + mock_prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) + + mock_cache = MagicMock() + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_cache.async_set_cache = AsyncMock() + + key_obj = await get_key_object( + hashed_token="hashed-token-1", + prisma_client=mock_prisma_client, + user_api_key_cache=mock_cache, + ) + + assert key_obj.token == "hashed-token-1" + assert mock_prisma_client.get_data.await_count == 2 + mock_prisma_client.attempt_db_reconnect.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_get_key_object_should_raise_if_reconnect_fails_on_db_connection_error(): + mock_prisma_client = MagicMock() + mock_prisma_client.get_data = AsyncMock( + side_effect=httpx.ConnectError("db not reachable after outage") + ) + mock_prisma_client.attempt_db_reconnect = AsyncMock(return_value=False) + + mock_cache = MagicMock() + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_cache.async_set_cache = AsyncMock() + + with pytest.raises(Exception, match="db not reachable after outage"): + await get_key_object( + hashed_token="hashed-token-2", + prisma_client=mock_prisma_client, + user_api_key_cache=mock_cache, + ) + + mock_prisma_client.attempt_db_reconnect.assert_awaited_once() + assert mock_prisma_client.get_data.await_count == 1 + + def test_get_cli_jwt_auth_token_default_expiration(valid_sso_user_defined_values): """Test generating CLI JWT token with default 24-hour expiration""" token = ExperimentalUIJWTToken.get_cli_jwt_auth_token(valid_sso_user_defined_values) @@ -180,9 +232,10 @@ def test_get_cli_jwt_auth_token_custom_expiration( ): """Test generating CLI JWT token with custom expiration via environment variable""" import importlib + from litellm import constants from litellm.proxy.auth import auth_checks - + # Set custom expiration to 48 hours monkeypatch.setenv("LITELLM_CLI_JWT_EXPIRATION_HOURS", "48")