Fix flaky test running blocking code in event loop (#155)

- Replace StepVerifier with assertWith for cleaner test assertions
- Add try-with-resources blocks for proper client resource management
- Use closeGracefully().block() for proper server shutdown

Signed-off-by: Dariusz Jędrzejczyk <dariusz.jedrzejczyk@broadcom.com>

mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/WebFluxSseIntegrationTests.java

@@ -37,10 +37,8 @@ import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.ValueSource;
-import reactor.core.publisher.Mono;
 import reactor.netty.DisposableServer;
 import reactor.netty.http.server.HttpServer;
-import reactor.test.StepVerifier;
 
 import org.springframework.http.server.reactive.HttpHandler;
 import org.springframework.http.server.reactive.ReactorHttpHandlerAdapter;
@@ -50,6 +48,7 @@ import org.springframework.web.reactive.function.server.RouterFunctions;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
+import static org.assertj.core.api.Assertions.assertWith;
 import static org.awaitility.Awaitility.await;
 import static org.mockito.Mockito.mock;
 
@@ -109,12 +108,9 @@ class WebFluxSseIntegrationTests {
 		var clientBuilder = clientBuilders.get(clientType);
 
 		McpServerFeatures.AsyncToolSpecification tool = new McpServerFeatures.AsyncToolSpecification(
-				new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema), (exchange, request) -> {
+				new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema),
-
+				(exchange, request) -> exchange.createMessage(mock(CreateMessageRequest.class))
-					exchange.createMessage(mock(McpSchema.CreateMessageRequest.class)).block();
+					.thenReturn(mock(CallToolResult.class)));
-
-					return Mono.just(mock(CallToolResult.class));
-				});
 
 		var server = McpServer.async(mcpServerTransportProvider).serverInfo("test-server", "1.0.0").tools(tool).build();
 
@@ -151,6 +147,8 @@ class WebFluxSseIntegrationTests {
 		CallToolResult callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")),
 				null);
 
+		AtomicReference<CreateMessageResult> samplingResult = new AtomicReference<>();
+
 		McpServerFeatures.AsyncToolSpecification tool = new McpServerFeatures.AsyncToolSpecification(
 				new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema), (exchange, request) -> {
 
@@ -165,16 +163,9 @@ class WebFluxSseIntegrationTests {
 							.build())
 						.build();
 
-					StepVerifier.create(exchange.createMessage(craeteMessageRequest)).consumeNextWith(result -> {
+					return exchange.createMessage(craeteMessageRequest)
-						assertThat(result).isNotNull();
+						.doOnNext(samplingResult::set)
-						assertThat(result.role()).isEqualTo(Role.USER);
+						.thenReturn(callResponse);
-						assertThat(result.content()).isInstanceOf(McpSchema.TextContent.class);
-						assertThat(((McpSchema.TextContent) result.content()).text()).isEqualTo("Test message");
-						assertThat(result.model()).isEqualTo("MockModelName");
-						assertThat(result.stopReason()).isEqualTo(CreateMessageResult.StopReason.STOP_SEQUENCE);
-					}).verifyComplete();
-
-					return Mono.just(callResponse);
 				});
 
 		var mcpServer = McpServer.async(mcpServerTransportProvider)
@@ -194,8 +185,17 @@ class WebFluxSseIntegrationTests {
 
 			assertThat(response).isNotNull();
 			assertThat(response).isEqualTo(callResponse);
+
+			assertWith(samplingResult.get(), result -> {
+				assertThat(result).isNotNull();
+				assertThat(result.role()).isEqualTo(Role.USER);
+				assertThat(result.content()).isInstanceOf(McpSchema.TextContent.class);
+				assertThat(((McpSchema.TextContent) result.content()).text()).isEqualTo("Test message");
+				assertThat(result.model()).isEqualTo("MockModelName");
+				assertThat(result.stopReason()).isEqualTo(CreateMessageResult.StopReason.STOP_SEQUENCE);
+			});
 		}
-		mcpServer.close();
+		mcpServer.closeGracefully().block();
 	}
 
 	@ParameterizedTest(name = "{0} : {displayName} ")
@@ -218,16 +218,13 @@ class WebFluxSseIntegrationTests {
 					CreateMessageResult.StopReason.STOP_SEQUENCE);
 		};
 
-		var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0"))
-			.capabilities(ClientCapabilities.builder().sampling().build())
-			.sampling(samplingHandler)
-			.build();
-
 		// Server
 
 		CallToolResult callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")),
 				null);
 
+		AtomicReference<CreateMessageResult> samplingResult = new AtomicReference<>();
+
 		McpServerFeatures.AsyncToolSpecification tool = new McpServerFeatures.AsyncToolSpecification(
 				new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema), (exchange, request) -> {
 
@@ -242,16 +239,9 @@ class WebFluxSseIntegrationTests {
 							.build())
 						.build();
 
-					StepVerifier.create(exchange.createMessage(craeteMessageRequest)).consumeNextWith(result -> {
+					return exchange.createMessage(craeteMessageRequest)
-						assertThat(result).isNotNull();
+						.doOnNext(samplingResult::set)
-						assertThat(result.role()).isEqualTo(Role.USER);
+						.thenReturn(callResponse);
-						assertThat(result.content()).isInstanceOf(McpSchema.TextContent.class);
-						assertThat(((McpSchema.TextContent) result.content()).text()).isEqualTo("Test message");
-						assertThat(result.model()).isEqualTo("MockModelName");
-						assertThat(result.stopReason()).isEqualTo(CreateMessageResult.StopReason.STOP_SEQUENCE);
-					}).verifyComplete();
-
-					return Mono.just(callResponse);
 				});
 
 		var mcpServer = McpServer.async(mcpServerTransportProvider)
@@ -260,16 +250,30 @@ class WebFluxSseIntegrationTests {
 			.tools(tool)
 			.build();
 
-		InitializeResult initResult = mcpClient.initialize();
+		try (var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0"))
-		assertThat(initResult).isNotNull();
+			.capabilities(ClientCapabilities.builder().sampling().build())
+			.sampling(samplingHandler)
+			.build()) {
 
-		CallToolResult response = mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of()));
+			InitializeResult initResult = mcpClient.initialize();
+			assertThat(initResult).isNotNull();
 
-		assertThat(response).isNotNull();
+			CallToolResult response = mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of()));
-		assertThat(response).isEqualTo(callResponse);
 
-		mcpClient.close();
+			assertThat(response).isNotNull();
-		mcpServer.close();
+			assertThat(response).isEqualTo(callResponse);
+
+			assertWith(samplingResult.get(), result -> {
+				assertThat(result).isNotNull();
+				assertThat(result.role()).isEqualTo(Role.USER);
+				assertThat(result.content()).isInstanceOf(McpSchema.TextContent.class);
+				assertThat(((McpSchema.TextContent) result.content()).text()).isEqualTo("Test message");
+				assertThat(result.model()).isEqualTo("MockModelName");
+				assertThat(result.stopReason()).isEqualTo(CreateMessageResult.StopReason.STOP_SEQUENCE);
+			});
+		}
+
+		mcpServer.closeGracefully().block();
 	}
 
 	@ParameterizedTest(name = "{0} : {displayName} ")
@@ -283,7 +287,7 @@ class WebFluxSseIntegrationTests {
 			assertThat(request.messages()).hasSize(1);
 			assertThat(request.messages().get(0).content()).isInstanceOf(McpSchema.TextContent.class);
 			try {
-				TimeUnit.SECONDS.sleep(3);
+				TimeUnit.SECONDS.sleep(2);
 			}
 			catch (InterruptedException e) {
 				throw new RuntimeException(e);
@@ -292,11 +296,6 @@ class WebFluxSseIntegrationTests {
 					CreateMessageResult.StopReason.STOP_SEQUENCE);
 		};
 
-		var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0"))
-			.capabilities(ClientCapabilities.builder().sampling().build())
-			.sampling(samplingHandler)
-			.build();
-
 		// Server
 
 		CallToolResult callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")),
@@ -308,24 +307,9 @@ class WebFluxSseIntegrationTests {
 					var craeteMessageRequest = McpSchema.CreateMessageRequest.builder()
 						.messages(List.of(new McpSchema.SamplingMessage(McpSchema.Role.USER,
 								new McpSchema.TextContent("Test message"))))
-						.modelPreferences(ModelPreferences.builder()
-							.hints(List.of())
-							.costPriority(1.0)
-							.speedPriority(1.0)
-							.intelligencePriority(1.0)
-							.build())
 						.build();
 
-					StepVerifier.create(exchange.createMessage(craeteMessageRequest)).consumeNextWith(result -> {
+					return exchange.createMessage(craeteMessageRequest).thenReturn(callResponse);
-						assertThat(result).isNotNull();
-						assertThat(result.role()).isEqualTo(Role.USER);
-						assertThat(result.content()).isInstanceOf(McpSchema.TextContent.class);
-						assertThat(((McpSchema.TextContent) result.content()).text()).isEqualTo("Test message");
-						assertThat(result.model()).isEqualTo("MockModelName");
-						assertThat(result.stopReason()).isEqualTo(CreateMessageResult.StopReason.STOP_SEQUENCE);
-					}).verifyComplete();
-
-					return Mono.just(callResponse);
 				});
 
 		var mcpServer = McpServer.async(mcpServerTransportProvider)
@@ -334,15 +318,21 @@ class WebFluxSseIntegrationTests {
 			.tools(tool)
 			.build();
 
-		InitializeResult initResult = mcpClient.initialize();
+		try (var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0"))
-		assertThat(initResult).isNotNull();
+			.capabilities(ClientCapabilities.builder().sampling().build())
+			.sampling(samplingHandler)
+			.build()) {
 
-		assertThatExceptionOfType(McpError.class).isThrownBy(() -> {
+			InitializeResult initResult = mcpClient.initialize();
-			mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of()));
+			assertThat(initResult).isNotNull();
-		}).withMessageContaining("Timeout");
 
-		mcpClient.close();
+			assertThatExceptionOfType(McpError.class).isThrownBy(() -> {
-		mcpServer.close();
+				mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of()));
+			}).withMessageContaining("within 1000ms");
+
+		}
+
+		mcpServer.closeGracefully().block();
 	}
 
 	// ---------------------------------------
@@ -412,9 +402,8 @@ class WebFluxSseIntegrationTests {
 		var mcpServer = McpServer.sync(mcpServerTransportProvider).rootsChangeHandler((exchange, rootsUpdate) -> {
 		}).tools(tool).build();
 
-		try (
+		// Create client without roots capability
-				// Create client without roots capability
+		try (var mcpClient = clientBuilder.capabilities(ClientCapabilities.builder().build()).build()) {
-				var mcpClient = clientBuilder.capabilities(ClientCapabilities.builder().build()).build()) {
 
 			assertThat(mcpClient.initialize()).isNotNull();