unit tested mock

Author: mikeedjones

tests/unittests/tools/mcp_tool/test_mcp_session_manager.py

@@ -144,16 +144,14 @@ def test_init_with_streamable_http_params(self):
 
     assert manager._connection_params == http_params
 
-  @pytest.mark.asyncio
-  async def test_init_with_streamable_http_custom_httpx_factory(self):
-    """Test initialization with StreamableHTTPConnectionParams."""
-    import httpx
-    custom_httpx_client = httpx.AsyncClient()
+  @patch("google.adk.tools.mcp_tool.mcp_session_manager.streamablehttp_client")
+  def test_init_with_streamable_http_custom_httpx_factory(
+    self, mock_streamablehttp_client
+  ):
+    """Test that streamablehttp_client is called with custom httpx_client_factory."""
+    from datetime import timedelta
 
-    def _httpx_factory(headers=None, timeout=None, auth=None):
-      return custom_httpx_client
-
-    custom_httpx_factory = Mock(side_effect=_httpx_factory)
+    custom_httpx_factory = Mock()
 
     http_params = StreamableHTTPConnectionParams(
       url="https://example.com/mcp",
@@ -162,9 +160,16 @@ def _httpx_factory(headers=None, timeout=None, auth=None):
     )
     manager = MCPSessionManager(http_params)
 
-    async with manager._create_client():
-      #assert factory was called
-      custom_httpx_factory.assert_called_once()
+    manager._create_client()
+
+    mock_streamablehttp_client.assert_called_once_with(
+      url="https://example.com/mcp",
+      headers=None,
+      timeout=timedelta(seconds=15.0),
+      sse_read_timeout=timedelta(seconds=300.0),
+      terminate_on_close=True,
+      httpx_client_factory=custom_httpx_factory,
+    )