From 5e0fbf079c00eefd2bb4b8c31baa53f918fb2d85 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Fri, 20 Feb 2026 15:53:15 -0800 Subject: [PATCH] test(auth): assert reconnect timeout argument in key lookup --- tests/test_litellm/proxy/auth/test_auth_checks.py | 10 ++++++++-- 1 file changed, 8 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 8aad1ebe684..4c2909384bd 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -177,7 +177,10 @@ async def test_get_key_object_should_reconnect_once_on_db_connection_error(): 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() + mock_prisma_client.attempt_db_reconnect.assert_awaited_once_with( + reason="auth_get_key_object_lookup_failure", + timeout_seconds=2.0, + ) @pytest.mark.asyncio @@ -199,7 +202,10 @@ async def test_get_key_object_should_raise_if_reconnect_fails_on_db_connection_e user_api_key_cache=mock_cache, ) - mock_prisma_client.attempt_db_reconnect.assert_awaited_once() + mock_prisma_client.attempt_db_reconnect.assert_awaited_once_with( + reason="auth_get_key_object_lookup_failure", + timeout_seconds=2.0, + ) assert mock_prisma_client.get_data.await_count == 1