Fix the simulateIncomingMessage coordination

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

Author: tzolov

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

@@ -4,6 +4,7 @@
 
 package io.modelcontextprotocol;
 
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Function;
 
@@ -33,37 +34,40 @@ public class MockMcpTransport implements ClientMcpTransport, ServerMcpTransport
 
 	private final Flux<McpSchema.JSONRPCMessage> outboundView = outgoing.asFlux().cache(1);
 
+	java.util.concurrent.CountDownLatch latch = new java.util.concurrent.CountDownLatch(1);
+
 	public void simulateIncomingMessage(McpSchema.JSONRPCMessage message) {
 		if (inbound.tryEmitNext(message).isFailure()) {
 			throw new RuntimeException("Failed to emit message " + message);
 		}
 		inboundMessageCount.incrementAndGet();
-
-		try {
-			Thread.sleep(200);
-		}
-		catch (InterruptedException e) {
-			e.printStackTrace();
-		}
+		latch = new java.util.concurrent.CountDownLatch(1);
 	}
 
 	@Override
 	public Mono<Void> sendMessage(McpSchema.JSONRPCMessage message) {
 		if (outgoing.tryEmitNext(message).isFailure()) {
 			return Mono.error(new RuntimeException("Can't emit outgoing message " + message));
 		}
+		latch.countDown();
 		return Mono.empty();
 	}
 
 	public McpSchema.JSONRPCRequest getLastSentMessageAsRequest() {
-		return (JSONRPCRequest) outboundView.blockFirst();
+		return (JSONRPCRequest) getLastSentMessage();
 	}
 
 	public McpSchema.JSONRPCNotification getLastSentMessageAsNotifiation() {
-		return (JSONRPCNotification) outboundView.blockFirst();
+		return (JSONRPCNotification) getLastSentMessage();
 	}
 
 	public McpSchema.JSONRPCMessage getLastSentMessage() {
+		try {
+			latch.await(200, TimeUnit.MILLISECONDS);
+		}
+		catch (InterruptedException e) {
+			e.printStackTrace();
+		}
 		return outboundView.blockFirst();
 	}
 

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

@@ -37,9 +37,6 @@ private InitializeResult initialization(McpAsyncClient asyncMcpClient, MockMcpTr
 					.build(),
 				new McpSchema.Implementation("test-server", "1.0.0"), "Test instructions");
 
-		// Use CountDownLatch to coordinate between threads
-		java.util.concurrent.CountDownLatch latch = new java.util.concurrent.CountDownLatch(1);
-
 		// Create a Mono that will handle the initialization and response simulation
 		return asyncMcpClient.initialize().doOnSubscribe(subscription -> {
 			// Run in a separate reactive context to avoid blocking the main subscription
@@ -51,17 +48,7 @@ private InitializeResult initialization(McpAsyncClient asyncMcpClient, MockMcpTr
 				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();
-		}).doOnTerminate(() -> {
-			try {
-				// Wait for the response simulation to complete
-				latch.await(5, java.util.concurrent.TimeUnit.SECONDS);
-			}
-			catch (InterruptedException e) {
-				Thread.currentThread().interrupt();
-				throw new RuntimeException("Interrupted while waiting for initialization", e);
-			}
 		}).block();
 	}
 
@@ -82,9 +69,6 @@ void testSuccessfulInitialization() {
 		McpSchema.InitializeResult mockInitResult = new McpSchema.InitializeResult(McpSchema.LATEST_PROTOCOL_VERSION,
 				mockServerCapabilities, mockServerInfo, "Test instructions");
 
-		// Use CountDownLatch to coordinate between threads
-		java.util.concurrent.CountDownLatch latch = new java.util.concurrent.CountDownLatch(1);
-
 		// Start initialization with reactive handling
 		InitializeResult result = asyncMcpClient.initialize().doOnSubscribe(subscription -> {
 			// Run in a separate reactive context to avoid blocking the main subscription
@@ -96,17 +80,8 @@ void testSuccessfulInitialization() {
 				McpSchema.JSONRPCResponse initResponse = new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION,
 						initRequest.id(), mockInitResult, null);
 				transport.simulateIncomingMessage(initResponse);
-				latch.countDown();
+				// latch.countDown();
 			}).subscribeOn(reactor.core.scheduler.Schedulers.boundedElastic()).subscribe();
-		}).doOnTerminate(() -> {
-			try {
-				// Wait for the response simulation to complete
-				latch.await(5, java.util.concurrent.TimeUnit.SECONDS);
-			}
-			catch (InterruptedException e) {
-				Thread.currentThread().interrupt();
-				throw new RuntimeException("Interrupted while waiting for initialization", e);
-			}
 		}).block();
 
 		// Verify initialized notification was sent