Rewrite testMissingHandlerReturnsMethodNotFoundError using a client transport

Mirrors the setup on HttpServletStreamableIntegrationTests.

Author: maff

mcp-test/src/test/java/io/modelcontextprotocol/server/HttpServletStatelessIntegrationTests.java

@@ -4,15 +4,12 @@
 
 package io.modelcontextprotocol.server;
 
-import java.net.URI;
-import java.net.http.HttpClient;
-import java.net.http.HttpRequest;
-import java.net.http.HttpResponse;
 import java.time.Duration;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.BiFunction;
+import java.util.function.Function;
 
 import io.modelcontextprotocol.client.McpClient;
 import io.modelcontextprotocol.client.transport.HttpClientStreamableHttpTransport;
@@ -49,6 +46,8 @@
 import org.springframework.mock.web.MockHttpServletRequest;
 import org.springframework.mock.web.MockHttpServletResponse;
 import org.springframework.web.client.RestClient;
+import reactor.core.publisher.Mono;
+import reactor.test.StepVerifier;
 
 import static io.modelcontextprotocol.server.transport.HttpServletStatelessServerTransport.APPLICATION_JSON;
 import static io.modelcontextprotocol.server.transport.HttpServletStatelessServerTransport.TEXT_EVENT_STREAM;
@@ -769,40 +768,46 @@ void testThrownMcpErrorAndJsonRpcError() throws Exception {
 		mcpServer.close();
 	}
 
-	@ParameterizedTest
-	@ValueSource(strings = { "tools/list", "resources/list", "prompts/list" })
-	void testMissingHandlerReturnsMethodNotFoundError(String method) throws Exception {
+	@Test
+	void testMissingHandlerReturnsMethodNotFoundError() {
 		var mcpServer = McpServer.sync(mcpStatelessServerTransport)
 			.serverInfo("test-server", "1.0.0")
 			.capabilities(ServerCapabilities.builder().build())
 			.build();
+		var clientTransport = HttpClientStreamableHttpTransport.builder("http://localhost:" + PORT)
+			.endpoint(CUSTOM_MESSAGE_ENDPOINT)
+			.build();
 
-		HttpResponse<String> response;
+		try (var mcpClient = McpClient.sync(clientTransport).build()) {
+			// Create a session using an MCP client
+			McpSchema.InitializeResult initResult = mcpClient.initialize();
+			assertThat(initResult).isNotNull();
 
-		try {
-			HttpRequest request = HttpRequest.newBuilder()
-				.uri(URI.create("http://localhost:" + PORT + CUSTOM_MESSAGE_ENDPOINT))
-				.header("Content-Type", "application/json")
-				.header("Accept", "application/json, text/event-stream")
-				.POST(HttpRequest.BodyPublishers.ofString("""
-						{
-							"jsonrpc": "2.0",
-							"method": "%s",
-							"id": "test-request-123",
-							"params": {}
+			// Override the response handler in the client to capture responses
+			AtomicReference<McpSchema.JSONRPCResponse> response = new AtomicReference<>();
+			var handler = (Function<Mono<McpSchema.JSONRPCMessage>, Mono<McpSchema.JSONRPCMessage>>) (
+					message) -> message.doOnNext(r -> {
+						if (r instanceof McpSchema.JSONRPCResponse resp) {
+							response.set(resp);
 						}
-						""".formatted(method)))
-				.build();
+					});
+			StepVerifier.create(clientTransport.connect(handler)).verifyComplete();
+
+			// Send a request for a non-existent method through the transport, bypassing
+			// the client's capability checks
+			StepVerifier
+				.create(clientTransport.sendMessage(new McpSchema.JSONRPCRequest("foo/bar", "test-request-123")))
+				.verifyComplete();
 
-			response = HttpClient.newHttpClient().send(request, HttpResponse.BodyHandlers.ofString());
+			// Wait until we've received the response
+			await().atMost(Duration.ofSeconds(1)).until(() -> response.get() != null);
+
+			assertThat(response.get().error().code()).isEqualTo(McpSchema.ErrorCodes.METHOD_NOT_FOUND);
+			assertThat(response.get().error().message()).isEqualTo("Method not found: foo/bar");
 		}
 		finally {
 			mcpServer.closeGracefully();
 		}
-
-		final var responseBody = response.body();
-		assertThatJson(responseBody).inPath("error.code").isEqualTo(McpSchema.ErrorCodes.METHOD_NOT_FOUND);
-		assertThatJson(responseBody).inPath("error.message").isEqualTo("Method not found: " + method);
 	}
 
 	private double evaluateExpression(String expression) {