Fix some unit tests

Signed-off-by: JermaineHua <[email protected]>

Author: CrazyHZM

mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/client/transport/WebFluxSseClientTransportProvider.java

@@ -263,39 +263,7 @@ public WebFluxSseClientTransport() {
 		 */
 		@Override
 		public Mono<Void> connect(Function<Mono<McpSchema.JSONRPCMessage>, Mono<McpSchema.JSONRPCMessage>> handler) {
-			Flux<ServerSentEvent<String>> events = eventStream();
-			inboundSubscription = events
-				.concatMap(event -> Mono.just(event).<McpSchema.JSONRPCMessage>handle((e, s) -> {
-					if (ENDPOINT_EVENT_TYPE.equals(event.event())) {
-						String messageEndpointUri = event.data();
-						if (messageEndpointSink.tryEmitValue(messageEndpointUri).isSuccess()) {
-							session.setId(event.id());
-							s.complete();
-						}
-						else {
-							// TODO: clarify with the spec if multiple events can be
-							// received
-							s.error(new McpError("Failed to handle SSE endpoint event"));
-						}
-					}
-					else if (MESSAGE_EVENT_TYPE.equals(event.event())) {
-						try {
-							McpSchema.JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(objectMapper,
-									event.data());
-							s.next(message);
-						}
-						catch (IOException ioException) {
-							s.error(ioException);
-						}
-					}
-					else {
-						s.error(new McpError("Received unrecognized SSE event type: " + event.event()));
-					}
-				}).flatMap(message -> session.handle(message)))
-				.subscribe();
-
-			// The connection is established once the server sends the endpoint event
-			return messageEndpointSink.asMono().then();
+			return connect();
 		}
 
 		@Override

mcp/src/main/java/io/modelcontextprotocol/client/transport/HttpClientSseClientTransportProvider.java

@@ -381,6 +381,11 @@ public HttpClientSseClientTransport(FlowSseClient sseClient, HttpClient httpClie
 		 */
 		@Override
 		public Mono<Void> connect(Function<Mono<JSONRPCMessage>, Mono<JSONRPCMessage>> handler) {
+			return connect();
+		}
+
+		@Override
+		public Mono<Void> connect() {
 			CompletableFuture<Void> future = new CompletableFuture<>();
 			connectionFuture.set(future);
 			sseClient.subscribe(baseUri + sseEndpoint, new FlowSseClient.SseEventHandler() {
@@ -394,7 +399,6 @@ public void onEvent(SseEvent event) {
 						if (ENDPOINT_EVENT_TYPE.equals(event.type())) {
 							String endpoint = event.data();
 							messageEndpoint.set(endpoint);
-							session.setId(getSessionIdFromUrl(endpoint));
 							closeLatch.countDown();
 							future.complete(null);
 						}

mcp/src/main/java/io/modelcontextprotocol/client/transport/StdioClientTransportProvider.java

@@ -179,38 +179,7 @@ public StdioClientTransport(ServerParameters params, ObjectMapper objectMapper)
 		 */
 		@Override
 		public Mono<Void> connect(Function<Mono<McpSchema.JSONRPCMessage>, Mono<McpSchema.JSONRPCMessage>> handler) {
-			return Mono.<Void>fromRunnable(() -> {
-				handleIncomingMessages(handler);
-				handleIncomingErrors();
-
-				// Prepare command and environment
-				List<String> fullCommand = new ArrayList<>();
-				fullCommand.add(params.getCommand());
-				fullCommand.addAll(params.getArgs());
-
-				ProcessBuilder processBuilder = this.getProcessBuilder();
-				processBuilder.command(fullCommand);
-				processBuilder.environment().putAll(params.getEnv());
-
-				// Start the process
-				try {
-					this.process = processBuilder.start();
-				}
-				catch (IOException e) {
-					throw new RuntimeException("Failed to start process with command: " + fullCommand, e);
-				}
-
-				// Validate process streams
-				if (this.process.getInputStream() == null || process.getOutputStream() == null) {
-					this.process.destroy();
-					throw new RuntimeException("Process input or output stream is null");
-				}
-
-				// Start threads
-				startInboundProcessing();
-				startOutboundProcessing();
-				startErrorProcessing();
-			}).subscribeOn(Schedulers.boundedElastic());
+			return connect();
 		}
 
 		@Override

mcp/src/test/java/io/modelcontextprotocol/client/StdioMcpSyncClientTests.java

@@ -5,12 +5,14 @@
 package io.modelcontextprotocol.client;
 
 import java.time.Duration;
+import java.util.Map;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicReference;
 
 import io.modelcontextprotocol.client.transport.ServerParameters;
 import io.modelcontextprotocol.client.transport.StdioClientTransportProvider;
+import io.modelcontextprotocol.spec.McpClientSession;
 import io.modelcontextprotocol.spec.McpClientTransport;
 import io.modelcontextprotocol.spec.McpClientTransportProvider;
 import org.junit.jupiter.api.Test;
@@ -49,6 +51,8 @@ void customErrorHandlerShouldReceiveErrors() throws InterruptedException {
 		AtomicReference<String> receivedError = new AtomicReference<>();
 
 		McpClientTransportProvider transportProvider = createMcpClientTransportProvider();
+		transportProvider.setSessionFactory(
+				(transport) -> new McpClientSession(Duration.ofSeconds(5), transport, Map.of(), Map.of()));
 		McpClientTransport transport = transportProvider.getSession().getTransport();
 		StepVerifier.create(transport.connect(msg -> msg)).verifyComplete();
 

mcp/src/test/java/io/modelcontextprotocol/client/transport/HttpClientSseClientTransportTests.java

@@ -124,7 +124,7 @@ void setUp() {
 		transportProvider.setSessionFactory(
 				(transport) -> new McpClientSession(Duration.ofSeconds(5), transport, Map.of(), Map.of()));
 		transport = new TestHttpClientSseClientTransport(transportProvider.getSession().getTransport());
-		transport.connect(Function.identity()).block();
+		transport.connect().block();
 	}
 
 	@AfterEach
@@ -334,6 +334,9 @@ void testCustomizeClient() {
 				customizerCalled.set(true);
 			})
 			.build();
+		customizedTransport.setSessionFactory(
+				(transport) -> new McpClientSession(Duration.ofSeconds(5), transport, Map.of(), Map.of()));
+		customizedTransport.getSession();
 
 		// Verify the customizer was called
 		assertThat(customizerCalled.get()).isTrue();
@@ -364,6 +367,9 @@ void testCustomizeRequest() {
 				headerValue.set(request.headers().firstValue("X-Custom-Header").orElse(null));
 			})
 			.build();
+		customizedTransport.setSessionFactory(
+				(transport) -> new McpClientSession(Duration.ofSeconds(5), transport, Map.of(), Map.of()));
+		customizedTransport.getSession();
 
 		// Verify the customizer was called
 		assertThat(customizerCalled.get()).isTrue();
@@ -393,6 +399,9 @@ void testChainedCustomizations() {
 				requestCustomizerCalled.set(true);
 			})
 			.build();
+		customizedTransport.setSessionFactory(
+				(transport) -> new McpClientSession(Duration.ofSeconds(5), transport, Map.of(), Map.of()));
+		customizedTransport.getSession();
 
 		// Verify both customizers were called
 		assertThat(clientCustomizerCalled.get()).isTrue();

mcp/src/test/java/io/modelcontextprotocol/spec/McpClientSessionTests.java

@@ -9,6 +9,7 @@
 
 import com.fasterxml.jackson.core.type.TypeReference;
 import io.modelcontextprotocol.MockMcpClientTransport;
+import io.modelcontextprotocol.MockMcpClientTransportProvider;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -41,13 +42,18 @@ class McpClientSessionTests {
 
 	private McpClientSession session;
 
-	private MockMcpClientTransport transport;
+	private MockMcpClientTransportProvider.MockMcpClientTransport transport;
+
+	private MockMcpClientTransportProvider transportProvider;
 
 	@BeforeEach
 	void setUp() {
-		transport = new MockMcpClientTransport();
-		session = new McpClientSession(TIMEOUT, transport, Map.of(),
-				Map.of(TEST_NOTIFICATION, params -> Mono.fromRunnable(() -> logger.info("Status update: " + params))));
+		transportProvider = new MockMcpClientTransportProvider();
+		transportProvider.setSessionFactory((transport) -> new McpClientSession(TIMEOUT, transport, Map.of(),
+				Map.of(TEST_NOTIFICATION, params -> Mono.fromRunnable(() -> logger.info("Status update: " + params)))));
+		session = transportProvider.getSession();
+		transport = transportProvider.getTransport();
+
 	}
 
 	@AfterEach
@@ -139,8 +145,11 @@ void testRequestHandling() {
 		String echoMessage = "Hello MCP!";
 		Map<String, McpClientSession.RequestHandler<?>> requestHandlers = Map.of(ECHO_METHOD,
 				params -> Mono.just(params));
-		transport = new MockMcpClientTransport();
-		session = new McpClientSession(TIMEOUT, transport, requestHandlers, Map.of());
+		transportProvider = new MockMcpClientTransportProvider();
+		transportProvider
+			.setSessionFactory((transport) -> new McpClientSession(TIMEOUT, transport, requestHandlers, Map.of()));
+		session = transportProvider.getSession();
+		transport = transportProvider.getTransport();
 
 		// Simulate incoming request
 		McpSchema.JSONRPCRequest request = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, ECHO_METHOD,
@@ -159,9 +168,11 @@ void testRequestHandling() {
 	void testNotificationHandling() {
 		Sinks.One<Object> receivedParams = Sinks.one();
 
-		transport = new MockMcpClientTransport();
-		session = new McpClientSession(TIMEOUT, transport, Map.of(),
-				Map.of(TEST_NOTIFICATION, params -> Mono.fromRunnable(() -> receivedParams.tryEmitValue(params))));
+		transportProvider = new MockMcpClientTransportProvider();
+		transportProvider.setSessionFactory((transport) -> new McpClientSession(TIMEOUT, transport, Map.of(),
+				Map.of(TEST_NOTIFICATION, params -> Mono.fromRunnable(() -> receivedParams.tryEmitValue(params)))));
+		session = transportProvider.getSession();
+		transport = transportProvider.getTransport();
 
 		// Simulate incoming notification from the server
 		Map<String, Object> notificationParams = Map.of("status", "ready");