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

refactor tests

Browse source
This commit is contained in:
Ishaan Jaff 2025-03-26 20:00:10 -07:00
parent 88ef97b9d1
commit 15c04da735
3 changed files with 110 additions and 72 deletions
Download patch file Download diff file Expand all files Collapse all files

72
tests/litellm/proxy/auth/test_auth_exception_handler.py
View file

@ -29,78 +29,6 @@ from litellm.proxy._types import ProxyErrorTypes, ProxyException
from litellm.proxy.auth.auth_exception_handler import UserAPIKeyAuthExceptionHandler
# Test is_database_connection_error method
@pytest.mark.parametrize(
"prisma_error",
[
PrismaError(),
DataError(data={"user_facing_error": {"meta": {"table": "test_table"}}}),
UniqueViolationError(
data={"user_facing_error": {"meta": {"table": "test_table"}}}
),
ForeignKeyViolationError(
data={"user_facing_error": {"meta": {"table": "test_table"}}}
),
MissingRequiredValueError(
data={"user_facing_error": {"meta": {"table": "test_table"}}}
),
RawQueryError(data={"user_facing_error": {"meta": {"table": "test_table"}}}),
TableNotFoundError(
data={"user_facing_error": {"meta": {"table": "test_table"}}}
),
RecordNotFoundError(
data={"user_facing_error": {"meta": {"table": "test_table"}}}
),
HTTPClientClosedError(),
ClientNotConnectedError(),
],
)
def test_is_database_connection_error_prisma_errors(prisma_error):
"""
Test that all Prisma errors are considered database connection errors
"""
handler = UserAPIKeyAuthExceptionHandler()
assert handler.is_database_connection_error(prisma_error) == True
def test_is_database_connection_generic_errors():
"""
Test non-Prisma error cases for database connection checking
"""
handler = UserAPIKeyAuthExceptionHandler()
# Test with ProxyException (DB connection)
db_proxy_exception = ProxyException(
message="DB Connection Error",
type=ProxyErrorTypes.no_db_connection,
param="test-param",
)
assert handler.is_database_connection_error(db_proxy_exception) == True
# Test with non-DB error
regular_exception = Exception("Regular error")
assert handler.is_database_connection_error(regular_exception) == False
# Test should_allow_request_on_db_unavailable method
@patch(
"litellm.proxy.proxy_server.general_settings",
{"allow_requests_on_db_unavailable": True},
)
def test_should_allow_request_on_db_unavailable_true():
handler = UserAPIKeyAuthExceptionHandler()
assert handler.should_allow_request_on_db_unavailable() == True
@patch(
"litellm.proxy.proxy_server.general_settings",
{"allow_requests_on_db_unavailable": False},
)
def test_should_allow_request_on_db_unavailable_false():
handler = UserAPIKeyAuthExceptionHandler()
assert handler.should_allow_request_on_db_unavailable() == False
@pytest.mark.asyncio
@pytest.mark.parametrize(
"prisma_error",