Skip to content

Commit a1fcb65

Browse files
committed
unit tested mock
1 parent ee16c0e commit a1fcb65

File tree

1 file changed

+17
-12
lines changed

1 file changed

+17
-12
lines changed

tests/unittests/tools/mcp_tool/test_mcp_session_manager.py

Lines changed: 17 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -144,16 +144,14 @@ def test_init_with_streamable_http_params(self):
144144

145145
assert manager._connection_params == http_params
146146

147-
@pytest.mark.asyncio
148-
async def test_init_with_streamable_http_custom_httpx_factory(self):
149-
"""Test initialization with StreamableHTTPConnectionParams."""
150-
import httpx
151-
custom_httpx_client = httpx.AsyncClient()
147+
@patch("google.adk.tools.mcp_tool.mcp_session_manager.streamablehttp_client")
148+
def test_init_with_streamable_http_custom_httpx_factory(
149+
self, mock_streamablehttp_client
150+
):
151+
"""Test that streamablehttp_client is called with custom httpx_client_factory."""
152+
from datetime import timedelta
152153

153-
def _httpx_factory(headers=None, timeout=None, auth=None):
154-
return custom_httpx_client
155-
156-
custom_httpx_factory = Mock(side_effect=_httpx_factory)
154+
custom_httpx_factory = Mock()
157155

158156
http_params = StreamableHTTPConnectionParams(
159157
url="https://example.com/mcp",
@@ -162,9 +160,16 @@ def _httpx_factory(headers=None, timeout=None, auth=None):
162160
)
163161
manager = MCPSessionManager(http_params)
164162

165-
async with manager._create_client():
166-
#assert factory was called
167-
custom_httpx_factory.assert_called_once()
163+
manager._create_client()
164+
165+
mock_streamablehttp_client.assert_called_once_with(
166+
url="https://example.com/mcp",
167+
headers=None,
168+
timeout=timedelta(seconds=15.0),
169+
sse_read_timeout=timedelta(seconds=300.0),
170+
terminate_on_close=True,
171+
httpx_client_factory=custom_httpx_factory,
172+
)
168173

169174

170175

0 commit comments

Comments
 (0)