Address review comments

Signed-off-by: Christian Tzolov <christian.tzolov@broadcom.com>

Author: tzolov

mcp/src/test/java/io/modelcontextprotocol/MockMcpTransport.java

@@ -4,17 +4,16 @@
 
 package io.modelcontextprotocol;
 
-import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Function;
 
 import com.fasterxml.jackson.core.type.TypeReference;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import io.modelcontextprotocol.spec.ClientMcpTransport;
 import io.modelcontextprotocol.spec.McpSchema;
-import io.modelcontextprotocol.spec.ServerMcpTransport;
 import io.modelcontextprotocol.spec.McpSchema.JSONRPCNotification;
 import io.modelcontextprotocol.spec.McpSchema.JSONRPCRequest;
+import io.modelcontextprotocol.spec.ServerMcpTransport;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.core.publisher.Sinks;
@@ -34,14 +33,20 @@ public class MockMcpTransport implements ClientMcpTransport, ServerMcpTransport
 
 	private final Flux<McpSchema.JSONRPCMessage> outboundView = outgoing.asFlux().cache(1);
 
+	// Latch to wait for the next message(s) to be sent in response of simulated incoming
+	// message
 	java.util.concurrent.CountDownLatch latch = new java.util.concurrent.CountDownLatch(1);
 
 	public void simulateIncomingMessage(McpSchema.JSONRPCMessage message) {
+		simulateIncomingMessage(message, 1);
+	}
+
+	public void simulateIncomingMessage(McpSchema.JSONRPCMessage message, int expectedResponseMessagesCount) {
 		if (inbound.tryEmitNext(message).isFailure()) {
 			throw new RuntimeException("Failed to emit message " + message);
 		}
 		inboundMessageCount.incrementAndGet();
-		latch = new java.util.concurrent.CountDownLatch(1);
+		latch = new java.util.concurrent.CountDownLatch(expectedResponseMessagesCount);
 	}
 
 	@Override
@@ -63,7 +68,7 @@ public McpSchema.JSONRPCNotification getLastSentMessageAsNotifiation() {
 
 	public McpSchema.JSONRPCMessage getLastSentMessage() {
 		try {
-			latch.await(200, TimeUnit.MILLISECONDS);
+			latch.await();
 		}
 		catch (InterruptedException e) {
 			e.printStackTrace();

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

@@ -60,29 +60,11 @@ void testSuccessfulInitialization() {
 		// Verify client is not initialized initially
 		assertThat(asyncMcpClient.isInitialized()).isFalse();
 
-		// Create mock server response
-		McpSchema.ServerCapabilities mockServerCapabilities = McpSchema.ServerCapabilities.builder()
-			.tools(true)
-			.resources(true, true) // Enable both resources and resource templates
-			.build();
-		McpSchema.Implementation mockServerInfo = new McpSchema.Implementation("test-server", "1.0.0");
-		McpSchema.InitializeResult mockInitResult = new McpSchema.InitializeResult(McpSchema.LATEST_PROTOCOL_VERSION,
-				mockServerCapabilities, mockServerInfo, "Test instructions");
-
 		// Start initialization with reactive handling
-		InitializeResult result = asyncMcpClient.initialize().doOnSubscribe(subscription -> {
-			// Run in a separate reactive context to avoid blocking the main subscription
-			Mono.fromRunnable(() -> {
-				McpSchema.JSONRPCRequest initRequest = transport.getLastSentMessageAsRequest();
-				assertThat(initRequest.method()).isEqualTo(McpSchema.METHOD_INITIALIZE);
+		InitializeResult result = initialization(asyncMcpClient, transport);
 
-				// Send mock server response
-				McpSchema.JSONRPCResponse initResponse = new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION,
-						initRequest.id(), mockInitResult, null);
-				transport.simulateIncomingMessage(initResponse);
-				// latch.countDown();
-			}).subscribeOn(reactor.core.scheduler.Schedulers.boundedElastic()).subscribe();
-		}).block();
+		// Verify client state after initialization
+		assertThat(asyncMcpClient.isInitialized()).isTrue();
 
 		// Verify initialized notification was sent
 		McpSchema.JSONRPCMessage notificationMessage = transport.getLastSentMessage();
@@ -93,15 +75,10 @@ void testSuccessfulInitialization() {
 		// Verify initialization result
 		assertThat(result).isNotNull();
 		assertThat(result.protocolVersion()).isEqualTo(McpSchema.LATEST_PROTOCOL_VERSION);
-		assertThat(result.capabilities()).isEqualTo(mockServerCapabilities);
-		assertThat(result.serverInfo()).isEqualTo(mockServerInfo);
+		assertThat(result.capabilities()).isEqualTo(asyncMcpClient.getServerCapabilities());
+		assertThat(result.serverInfo()).isEqualTo(asyncMcpClient.getServerInfo());
 		assertThat(result.instructions()).isEqualTo("Test instructions");
 
-		// Verify client state after initialization
-		assertThat(asyncMcpClient.isInitialized()).isTrue();
-		assertThat(asyncMcpClient.getServerCapabilities()).isEqualTo(mockServerCapabilities);
-		assertThat(asyncMcpClient.getServerInfo()).isEqualTo(mockServerInfo);
-
 		asyncMcpClient.closeGracefully();
 	}
 

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

@@ -82,7 +82,7 @@ void testSendRequest() {
 		StepVerifier.create(responseMono).then(() -> {
 			McpSchema.JSONRPCRequest request = transport.getLastSentMessageAsRequest();
 			transport.simulateIncomingMessage(
-					new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), responseData, null));
+					new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), responseData, null), 0);
 		}).consumeNextWith(response -> {
 			// Verify the request was sent
 			McpSchema.JSONRPCMessage sentMessage = transport.getLastSentMessageAsRequest();