diff --git a/tests/unit/integrations/test_custom_prompt_management.py b/tests/unit/integrations/test_custom_prompt_management.py index 0bf2063a98d..0255e12acad 100644 --- a/tests/unit/integrations/test_custom_prompt_management.py +++ b/tests/unit/integrations/test_custom_prompt_management.py @@ -55,6 +55,22 @@ class TestCustomPromptManagement(CustomPromptManagement): return model, messages, non_default_params +def _anthropic_chat_response() -> httpx.Response: + return httpx.Response( + 200, + json={ + "id": "msg_123", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "Hello"}], + "model": "claude-3-5-sonnet", + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ) + + @pytest.mark.asyncio async def test_custom_prompt_management_with_prompt_id(monkeypatch): custom_prompt_management = TestCustomPromptManagement() @@ -62,7 +78,7 @@ async def test_custom_prompt_management_with_prompt_id(monkeypatch): # Mock AsyncHTTPHandler.post method client = AsyncHTTPHandler() - with patch.object(client, "post", return_value=MagicMock()) as mock_post: + with patch.object(client, "post", return_value=_anthropic_chat_response()) as mock_post: await litellm.acompletion( model="anthropic/claude-3-5-sonnet", messages=[{"role": "user", "content": "Hello, how are you?"}], @@ -90,7 +106,7 @@ async def test_custom_prompt_management_with_prompt_id_and_prompt_variables(): # Mock AsyncHTTPHandler.post method client = AsyncHTTPHandler() - with patch.object(client, "post", return_value=MagicMock()) as mock_post: + with patch.object(client, "post", return_value=_anthropic_chat_response()) as mock_post: await litellm.acompletion( model="anthropic/claude-3-5-sonnet", messages=[], @@ -119,7 +135,7 @@ async def test_custom_prompt_management_without_prompt_id(): # Mock AsyncHTTPHandler.post method client = AsyncHTTPHandler() - with patch.object(client, "post", return_value=MagicMock()) as mock_post: + with patch.object(client, "post", return_value=_anthropic_chat_response()) as mock_post: await litellm.acompletion( model="anthropic/claude-3-5-sonnet", messages=[{"role": "user", "content": "Hello, how are you?"}], diff --git a/tests/unit/router_utils/test_pattern_match_deployments.py b/tests/unit/router_utils/test_pattern_match_deployments.py index 01af021112b..227df270675 100644 --- a/tests/unit/router_utils/test_pattern_match_deployments.py +++ b/tests/unit/router_utils/test_pattern_match_deployments.py @@ -317,7 +317,23 @@ def test_router_pattern_match_e2e(): ] ) - with patch.object(client, "post", new=MagicMock()) as mock_post: + with patch.object( + client, + "post", + return_value=httpx.Response( + 200, + json={ + "id": "msg_123", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "Hello"}], + "model": "claude-3-5-sonnet", + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ), + ) as mock_post: router.completion( model="llmengine/my-custom-model",