diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index a909c510581..4d2285b1971 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -15,9 +15,10 @@ sys.path.insert( ) # Adds the parent directory to the system path -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch from litellm.proxy.utils import get_custom_url, join_paths +from litellm.types.utils import CallTypes def test_get_custom_url(monkeypatch): @@ -71,6 +72,28 @@ def test_proxy_only_error_false_for_other_error_type(): ) +@pytest.mark.asyncio +async def test_pre_call_hook_does_not_track_mcp_tool_as_hanging_request(): + proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) + proxy_logging_obj.slack_alerting_instance = MagicMock() + proxy_logging_obj.slack_alerting_instance.alerting = ["slack"] + proxy_logging_obj.slack_alerting_instance.response_taking_too_long = AsyncMock() + + await proxy_logging_obj.pre_call_hook( + user_api_key_dict=MagicMock(), + data=None, + call_type=CallTypes.call_mcp_tool.value, + ) + proxy_logging_obj.slack_alerting_instance.response_taking_too_long.assert_not_called() + + await proxy_logging_obj.pre_call_hook( + user_api_key_dict=MagicMock(), + data=None, + call_type=CallTypes.acompletion.value, + ) + proxy_logging_obj.slack_alerting_instance.response_taking_too_long.assert_called_once_with(request_data=None) + + @pytest.mark.asyncio async def test_proxy_only_error_log_marks_no_upstream_llm_call(): """A proxy-gate error (auth/rate-limit) synthesizes a ``Logging`` object and diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py b/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py index dc702db2fdb..05005dae797 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py @@ -12,7 +12,6 @@ import litellm from litellm.exceptions import RejectedRequestError from litellm.integrations.custom_logger import CustomLogger from litellm.proxy.utils import ProxyLogging -from litellm.types.utils import CallTypes @pytest.fixture(autouse=True) @@ -103,25 +102,6 @@ async def test_pre_call_hook_returns_none_for_none_data(proxy_logging, make_user assert out is None -@pytest.mark.asyncio -async def test_pre_call_hook_does_not_track_mcp_tool_as_hanging_request( - proxy_logging, make_user_api_key_auth, mock_callbacks_disabled -): - data = {"messages": [{"role": "user", "content": "Tool: search"}], "model": "mcp-tool-call"} - proxy_logging.slack_alerting_instance = MagicMock() - proxy_logging.slack_alerting_instance.alerting = ["slack"] - proxy_logging.slack_alerting_instance.response_taking_too_long = AsyncMock() - - out = await proxy_logging.pre_call_hook( - user_api_key_dict=make_user_api_key_auth(), - data=data, - call_type=CallTypes.call_mcp_tool.value, - ) - - assert out is data - proxy_logging.slack_alerting_instance.response_taking_too_long.assert_not_called() - - @pytest.mark.asyncio async def test_pre_call_hook_invokes_pre_call_override(proxy_logging, make_user_api_key_auth, monkeypatch): captured: Dict[str, Any] = {}