test: add tests for uncovered branches in streamable HTTP tests

Author: perhapzz

tests/shared/test_streamable_http.py

@@ -1020,13 +1020,6 @@ def test_get_validation(basic_server: None, basic_server_url: str):
 
 
 # Client-specific fixtures
-@pytest.fixture
-async def http_client(basic_server: None, basic_server_url: str):
-    """Create test client matching the SSE test pattern."""
-    async with httpx.AsyncClient(base_url=basic_server_url) as client:
-        yield client
-
-
 @pytest.fixture
 async def initialized_client_session(basic_server: None, basic_server_url: str):
     """Create initialized StreamableHTTP client session."""
@@ -2293,3 +2286,189 @@ async def test_streamable_http_client_preserves_custom_with_mcp_headers(
 
                 assert "content-type" in headers_data
                 assert headers_data["content-type"] == "application/json"
+
+
+@pytest.mark.anyio
+async def test_replay_events_after_nonexistent_event_id():
+    """Test replay_events_after returns None for non-existent event ID."""
+    store = SimpleEventStore()
+
+    # Store some events first
+    stream_id = "stream-1"
+    await store.store_event(stream_id, types.JSONRPCResponse(jsonrpc="2.0", id="1", result={"key": "value"}))
+
+    # Try to replay after a non-existent event ID
+    callback = MagicMock()
+    result = await store.replay_events_after("999", callback)
+    assert result is None
+    callback.assert_not_called()
+
+
+@pytest.mark.anyio
+async def test_replay_events_after_replays_messages():
+    """Test replay_events_after correctly replays messages after a given event ID."""
+    store = SimpleEventStore()
+
+    stream_id = "stream-1"
+    msg1: types.JSONRPCMessage = types.JSONRPCResponse(jsonrpc="2.0", id="1", result={"first": True})
+    msg2: types.JSONRPCMessage = types.JSONRPCResponse(jsonrpc="2.0", id="2", result={"second": True})
+    # Store: priming event (None), then two real messages
+    eid0 = await store.store_event(stream_id, None)
+    eid1 = await store.store_event(stream_id, msg1)
+    eid2 = await store.store_event(stream_id, msg2)
+
+    # Replay after priming event — should get both messages
+    replayed: list[EventMessage] = []
+
+    async def callback(event_msg: EventMessage) -> None:
+        replayed.append(event_msg)
+
+    result = await store.replay_events_after(eid0, callback)
+    assert result == stream_id
+    assert len(replayed) == 2
+    assert replayed[0].event_id == eid1
+    assert replayed[1].event_id == eid2
+
+    # Replay after first message — should get only second
+    replayed.clear()
+    result = await store.replay_events_after(eid1, callback)
+    assert result == stream_id
+    assert len(replayed) == 1
+    assert replayed[0].event_id == eid2
+
+
+@pytest.mark.anyio
+async def test_streamable_http_client_slow_resource(initialized_client_session: ClientSession):
+    """Test reading a slow:// resource."""
+    result = await initialized_client_session.read_resource("slow://test-host")
+    assert len(result.contents) == 1
+    assert isinstance(result.contents[0], TextResourceContents)
+    assert result.contents[0].text == "Slow response from test-host"
+
+
+@pytest.mark.anyio
+async def test_streamable_http_client_long_running_with_checkpoints(basic_server: None, basic_server_url: str):
+    """Test calling the long_running_with_checkpoints tool."""
+    captured_notifications: list[str] = []
+
+    async def message_handler(
+        message: RequestResponder[types.ServerRequest, types.ClientResult] | types.ServerNotification | Exception,
+    ) -> None:
+        if isinstance(message, Exception):
+            return  # pragma: no cover
+        if isinstance(message, types.ServerNotification):  # pragma: no branch
+            if isinstance(message, types.LoggingMessageNotification):  # pragma: no branch
+                captured_notifications.append(str(message.params.data))
+
+    async with streamable_http_client(f"{basic_server_url}/mcp") as (read_stream, write_stream):
+        async with ClientSession(read_stream, write_stream, message_handler=message_handler) as session:
+            await session.initialize()
+
+            result = await session.call_tool("long_running_with_checkpoints", {})
+            assert len(result.content) == 1
+            assert result.content[0].type == "text"
+            assert isinstance(result.content[0], TextContent)
+            assert result.content[0].text == "Completed!"
+
+            # Should have received the two log notifications
+            assert "Tool started" in captured_notifications
+            assert "Tool is almost done" in captured_notifications
+
+
+@pytest.mark.anyio
+async def test_streamablehttp_server_sampling_non_text_content(basic_server: None, basic_server_url: str):
+    """Test server-initiated sampling where callback returns non-text content."""
+
+    async def sampling_callback(
+        context: RequestContext[ClientSession],
+        params: types.CreateMessageRequestParams,
+    ) -> types.CreateMessageResult:
+        return types.CreateMessageResult(
+            role="assistant",
+            content=types.ImageContent(
+                type="image",
+                data="base64data",
+                mime_type="image/png",
+            ),
+            model="test-model",
+            stop_reason="endTurn",
+        )
+
+    async with streamable_http_client(f"{basic_server_url}/mcp") as (read_stream, write_stream):
+        async with ClientSession(read_stream, write_stream, sampling_callback=sampling_callback) as session:
+            await session.initialize()
+
+            tool_result = await session.call_tool("test_sampling_tool", {})
+            assert len(tool_result.content) == 1
+            assert tool_result.content[0].type == "text"
+            # Non-text content should be stringified
+            assert "Response from sampling:" in tool_result.content[0].text
+
+
+@pytest.mark.anyio
+async def test_tool_with_multiple_stream_closes(
+    event_server: tuple[SimpleEventStore, str],
+) -> None:
+    """Test tool_with_multiple_stream_closes which calls close_sse_stream multiple times."""
+    _, server_url = event_server
+    captured_notifications: list[str] = []
+
+    async def message_handler(
+        message: RequestResponder[types.ServerRequest, types.ClientResult] | types.ServerNotification | Exception,
+    ) -> None:
+        if isinstance(message, Exception):
+            return  # pragma: no cover
+        if isinstance(message, types.ServerNotification):  # pragma: no branch
+            if isinstance(message, types.LoggingMessageNotification):  # pragma: no branch
+                captured_notifications.append(str(message.params.data))
+
+    async with streamable_http_client(f"{server_url}/mcp") as (read_stream, write_stream):
+        async with ClientSession(read_stream, write_stream, message_handler=message_handler) as session:
+            await session.initialize()
+
+            result = await session.call_tool(
+                "tool_with_multiple_stream_closes",
+                {"checkpoints": 3, "sleep_time": 0.2},
+            )
+            assert result.content[0].type == "text"
+            assert isinstance(result.content[0], TextContent)
+            assert "Completed 3 checkpoints" in result.content[0].text
+
+            # All checkpoint notifications should have been received
+            for i in range(3):
+                assert f"checkpoint_{i}" in captured_notifications
+
+
+@pytest.mark.anyio
+async def test_tool_with_multiple_stream_closes_no_event_store(
+    basic_server: None,
+    basic_server_url: str,
+) -> None:
+    """Test multi_close_tool without event store — close_sse_stream is None."""
+    async with streamable_http_client(f"{basic_server_url}/mcp") as (read_stream, write_stream):
+        async with ClientSession(read_stream, write_stream) as session:
+            await session.initialize()
+
+            result = await session.call_tool(
+                "tool_with_multiple_stream_closes",
+                {"checkpoints": 2, "sleep_time": 0.1},
+            )
+            assert result.content[0].type == "text"
+            assert isinstance(result.content[0], TextContent)
+            assert "Completed 2 checkpoints" in result.content[0].text
+
+
+@pytest.mark.anyio
+async def test_tool_with_standalone_stream_close_no_event_store(
+    basic_server: None,
+    basic_server_url: str,
+) -> None:
+    """Test standalone_stream_close without event store — close_standalone_sse_stream is None."""
+    async with streamable_http_client(f"{basic_server_url}/mcp") as (read_stream, write_stream):
+        async with ClientSession(read_stream, write_stream) as session:
+            await session.initialize()
+
+            result = await session.call_tool("tool_with_standalone_stream_close", {})
+            assert result.content[0].type == "text"
+            assert isinstance(result.content[0], TextContent)
+            assert result.content[0].text == "Standalone stream close test done"