import asyncio import json import os import sys from unittest.mock import MagicMock, patch import pytest from fastapi import HTTPException, Request, status from prisma import errors as prisma_errors from prisma.errors import ( ClientNotConnectedError, DataError, ForeignKeyViolationError, HTTPClientClosedError, MissingRequiredValueError, PrismaError, RawQueryError, RecordNotFoundError, TableNotFoundError, UniqueViolationError, ) sys.path.insert( 0, os.path.abspath("../../..") ) # Adds the parent directory to the system path from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ProxyErrorTypes, ProxyException from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler # 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 """ assert PrismaDBExceptionHandler.is_database_connection_error(prisma_error) == True def test_is_database_connection_generic_errors(): """ Test non-Prisma error cases for database connection checking """ assert ( PrismaDBExceptionHandler.is_database_connection_error( Exception("Regular error") ) == False ) # Test with ProxyException (DB connection) db_proxy_exception = ProxyException( message="DB Connection Error", type=ProxyErrorTypes.no_db_connection, param="test-param", ) assert ( PrismaDBExceptionHandler.is_database_connection_error(db_proxy_exception) == True ) # Test with non-DB error regular_exception = Exception("Regular error") assert ( PrismaDBExceptionHandler.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(): assert PrismaDBExceptionHandler.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(): assert PrismaDBExceptionHandler.should_allow_request_on_db_unavailable() == False