Skip to content

Commit 2e8ecdc

Browse files
Simplify test_session_group.py mocking
Remove complex mocking of create_mcp_http_client in the streamablehttp test case. Instead, let the real create_mcp_http_client execute and only verify that streamable_http_client receives the correct parameters including a real httpx.AsyncClient instance. This simplifies the test by: - Removing 13 lines of mock setup code - Removing 14 lines of mock verification code - Removing 3 lines of mock cleanup code - Trusting that create_mcp_http_client works (it has its own tests) The test now focuses on verifying the integration between session_group and streamable_http_client rather than re-testing create_mcp_http_client.
1 parent 95304e4 commit 2e8ecdc

File tree

1 file changed

+6
-34
lines changed

1 file changed

+6
-34
lines changed

tests/client/test_session_group.py

Lines changed: 6 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -292,20 +292,6 @@ async def test_establish_session_parameterized(
292292
):
293293
with mock.patch("mcp.client.session_group.mcp.ClientSession") as mock_ClientSession_class:
294294
with mock.patch(patch_target_for_client_func) as mock_specific_client_func:
295-
# For streamablehttp, also need to mock create_mcp_http_client
296-
if client_type_name == "streamablehttp":
297-
mock_create_http_client = mock.patch("mcp.client.session_group.create_mcp_http_client")
298-
mock_create_http_client_func = mock_create_http_client.start()
299-
# Mock httpx_client returned by create_mcp_http_client
300-
mock_httpx_client = mock.AsyncMock(name="MockHttpxClient")
301-
mock_httpx_client.__aenter__.return_value = mock_httpx_client
302-
mock_httpx_client.__aexit__ = mock.AsyncMock(return_value=None)
303-
mock_create_http_client_func.return_value = mock_httpx_client
304-
else:
305-
mock_create_http_client = None
306-
mock_create_http_client_func = None
307-
mock_httpx_client = None
308-
309295
mock_client_cm_instance = mock.AsyncMock(name=f"{client_type_name}ClientCM")
310296
mock_read_stream = mock.AsyncMock(name=f"{client_type_name}Read")
311297
mock_write_stream = mock.AsyncMock(name=f"{client_type_name}Write")
@@ -368,24 +354,14 @@ async def test_establish_session_parameterized(
368354
)
369355
elif client_type_name == "streamablehttp": # pragma: no branch
370356
assert isinstance(server_params_instance, StreamableHttpParameters)
371-
# Verify create_mcp_http_client was called with headers and timeout
357+
# Verify streamable_http_client was called with url, httpx_client, and terminate_on_close
358+
# The httpx_client is created by the real create_mcp_http_client
372359
import httpx
373360

374-
assert mock_create_http_client_func is not None
375-
expected_timeout = httpx.Timeout(
376-
server_params_instance.timeout.total_seconds(),
377-
read=server_params_instance.sse_read_timeout.total_seconds(),
378-
)
379-
mock_create_http_client_func.assert_called_once_with(
380-
headers=server_params_instance.headers,
381-
timeout=expected_timeout,
382-
)
383-
# Verify streamable_http_client was called with url, httpx_client, and terminate_on_close
384-
mock_specific_client_func.assert_called_once_with(
385-
url=server_params_instance.url,
386-
httpx_client=mock_httpx_client,
387-
terminate_on_close=server_params_instance.terminate_on_close,
388-
)
361+
call_args = mock_specific_client_func.call_args
362+
assert call_args.kwargs["url"] == server_params_instance.url
363+
assert call_args.kwargs["terminate_on_close"] == server_params_instance.terminate_on_close
364+
assert isinstance(call_args.kwargs["httpx_client"], httpx.AsyncClient)
389365

390366
mock_client_cm_instance.__aenter__.assert_awaited_once()
391367

@@ -407,7 +383,3 @@ async def test_establish_session_parameterized(
407383
# 3. Assert returned values
408384
assert returned_server_info is mock_initialize_result.serverInfo
409385
assert returned_session is mock_entered_session
410-
411-
# Clean up streamablehttp-specific mock
412-
if mock_create_http_client:
413-
mock_create_http_client.stop()

0 commit comments

Comments
 (0)