diff --git a/tests/llm_translation/test_cohere.py b/tests/llm_translation/test_cohere.py index 41dc491d35..76e94ed959 100644 --- a/tests/llm_translation/test_cohere.py +++ b/tests/llm_translation/test_cohere.py @@ -254,10 +254,17 @@ async def test_cohere_request_body_with_allowed_params(): } }] - client = AsyncHTTPHandler() + # Create a mock response + mock_response = AsyncMock() + mock_response.status_code = 200 + mock_response.json.return_value = { + "text": "I am Command, a language model developed by Cohere.", + "generation_id": "mock-generation-id", + "finish_reason": "COMPLETE" + } - # Mock the post method - with patch.object(client, "post", new=AsyncMock()) as mock_post: + # Mock the AsyncHTTPHandler.post method at the module level + with patch("litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", return_value=mock_response) as mock_post: try: await litellm.acompletion( model="cohere/command", @@ -265,8 +272,7 @@ async def test_cohere_request_body_with_allowed_params(): allowed_openai_params=["tools", "response_format", "reasoning_effort"], response_format=test_response_format, reasoning_effort=test_reasoning_effort, - tools=test_tools, - client=client + tools=test_tools ) except Exception: pass # We only care about the request body validation