Fix some unit tests

Signed-off-by: JermaineHua <crazyhzm@apache.org>

Author: CrazyHZM

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/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");