test docker model runner

tests/test_litellm/llms/docker_model_runner/test_docker_model_runner_chat_transformation.py

@@ -31,8 +31,13 @@ class TestDockerModelRunnerIntegration:
         1. Hits the correct URL: {api_base}/v1/chat/completions where api_base includes engine path
         2. Sends the correct request body with messages and parameters
         """
-        with patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") as mock_post:
-            # Mock the response
+        # Mock _get_httpx_client to return a mock HTTPHandler
+        with patch("litellm.llms.custom_httpx.llm_http_handler._get_httpx_client") as mock_get_client:
+            # Create a mock HTTPHandler instance
+            mock_http_handler = Mock()
+            mock_get_client.return_value = mock_http_handler
+            
+            # Create mock response
             mock_response = Mock(spec=httpx.Response)
             mock_response.json.return_value = {
                 "id": "chatcmpl-123",
@@ -55,7 +60,15 @@ class TestDockerModelRunnerIntegration:
             }
             mock_response.status_code = 200
             mock_response.headers = httpx.Headers({"content-type": "application/json"})
-            mock_post.return_value = mock_response
+            mock_response.text = json.dumps(mock_response.json.return_value)
+            
+            # Capture the request
+            captured_kwargs = {}
+            def mock_post(**kwargs):
+                captured_kwargs.update(kwargs)
+                return mock_response
+            
+            mock_http_handler.post.side_effect = mock_post
 
             # Make the completion call with engine in api_base
             response = completion(
@@ -66,21 +79,20 @@ class TestDockerModelRunnerIntegration:
                 max_tokens=100
             )
 
-            # Verify the URL was correct
-            assert mock_post.called
-            call_args = mock_post.call_args
-            url = call_args[1]["url"]
+            # Verify the request was made
+            assert mock_http_handler.post.called
+            url = captured_kwargs.get('url', '')
+            data = captured_kwargs.get('data', '')
             print("URL For request", url)
-            print("request body for request", json.dumps(call_args[1]["data"], indent=4))
+            print("request body for request", data)
             
             # Should hit {api_base}/v1/chat/completions where api_base includes engine
             assert "/engines/llama.cpp/v1/chat/completions" in url
             assert "http://localhost:22088" in url
 
             # Verify the request body
-            request_data = call_args[1]["data"]
-            if isinstance(request_data, str):
-                request_data = json.loads(request_data)
+            request_data = json.loads(data) if isinstance(data, str) else data
+            print("Parsed request data:", json.dumps(request_data, indent=4))
             
             # Check messages
             assert "messages" in request_data
@@ -102,8 +114,13 @@ class TestDockerModelRunnerIntegration:
         2. Specifies a different engine in the api_base
         3. Model name is sent in the request body
         """
-        with patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") as mock_post:
-            # Mock the response
+        # Mock _get_httpx_client to return a mock HTTPHandler
+        with patch("litellm.llms.custom_httpx.llm_http_handler._get_httpx_client") as mock_get_client:
+            # Create a mock HTTPHandler instance
+            mock_http_handler = Mock()
+            mock_get_client.return_value = mock_http_handler
+            
+            # Create mock response
             mock_response = Mock(spec=httpx.Response)
             mock_response.json.return_value = {
                 "id": "chatcmpl-456",
@@ -126,7 +143,15 @@ class TestDockerModelRunnerIntegration:
             }
             mock_response.status_code = 200
             mock_response.headers = httpx.Headers({"content-type": "application/json"})
-            mock_post.return_value = mock_response
+            mock_response.text = json.dumps(mock_response.json.return_value)
+            
+            # Capture the request
+            captured_kwargs = {}
+            def mock_post(**kwargs):
+                captured_kwargs.update(kwargs)
+                return mock_response
+            
+            mock_http_handler.post.side_effect = mock_post
 
             # Make the completion call with custom engine and host
             response = completion(
@@ -137,21 +162,20 @@ class TestDockerModelRunnerIntegration:
                 max_tokens=200
             )
 
-            # Verify the URL was correct
-            assert mock_post.called
-            call_args = mock_post.call_args
-            url = call_args[1]["url"]
+            # Verify the request was made
+            assert mock_http_handler.post.called
+            url = captured_kwargs.get('url', '')
+            data = captured_kwargs.get('data', '')
             print("URL For request", url)
-            print("request body for request", json.dumps(call_args[1]["data"], indent=4))
+            print("request body for request", data)
             
             # Should hit the custom host and engine
             assert "model-runner.docker.internal" in url
             assert "/engines/custom-engine/v1/chat/completions" in url
 
             # Verify the request body contains the model name
-            request_data = call_args[1]["data"]
-            if isinstance(request_data, str):
-                request_data = json.loads(request_data)
+            request_data = json.loads(data) if isinstance(data, str) else data
+            print("Parsed request data:", json.dumps(request_data, indent=4))
             
             # Check that model name is in the request body
             assert request_data["model"] == "mistral-7b"