refactor: improve integration tests stability

- Replace inline StepVerifier assertions with AtomicReference pattern for better testability
- Add assertWith() usage for cleaner assertion blocks in async tests
- Simplify reactive chains using doOnNext() and thenReturn() patterns
- Remove unnecessary Thread.sleep() from elicitation handler
- Improve variable naming (rootsRef → toolsRef) for clarity
- Update timeout error message assertions to be more specific
- Clean up code formatting and remove redundant comments

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

mcp-spring/mcp-spring-webmvc/src/test/java/io/modelcontextprotocol/server/WebMvcStreamableIntegrationTests.java

@@ -11,8 +11,6 @@ import org.apache.catalina.LifecycleException;
 import org.apache.catalina.LifecycleState;
 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 org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.web.reactive.function.client.WebClient;
@@ -29,7 +27,6 @@ import io.modelcontextprotocol.client.transport.WebClientStreamableHttpTransport
 import io.modelcontextprotocol.server.McpServer.AsyncSpecification;
 import io.modelcontextprotocol.server.McpServer.SyncSpecification;
 import io.modelcontextprotocol.server.transport.WebMvcStreamableServerTransportProvider;
-import io.modelcontextprotocol.spec.McpSchema;
 import reactor.core.scheduler.Schedulers;
 
 class WebMvcStreamableIntegrationTests extends AbstractMcpClientServerIntegrationTests {

mcp-test/src/main/java/io/modelcontextprotocol/AbstractMcpClientServerIntegrationTests.java

@@ -8,6 +8,7 @@ import static net.javacrumbs.jsonunit.assertj.JsonAssertions.assertThatJson;
 import static net.javacrumbs.jsonunit.assertj.JsonAssertions.json;
 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;
 
@@ -144,6 +145,8 @@ public abstract class AbstractMcpClientServerIntegrationTests {
 		CallToolResult callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")),
 				null);
 
+		AtomicReference<CreateMessageResult> samplingResult = new AtomicReference<>();
+
 		McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder()
 			.tool(Tool.builder().name("tool1").description("tool1 description").inputSchema(emptyJsonSchema).build())
 			.callHandler((exchange, request) -> {
@@ -159,37 +162,35 @@ public abstract class AbstractMcpClientServerIntegrationTests {
 						.build())
 					.build();
 
-				StepVerifier.create(exchange.createMessage(createMessageRequest)).consumeNextWith(result -> {
+				return exchange.createMessage(createMessageRequest)
-					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);
 			})
 			.build();
 
-		//@formatter:off		
+		var mcpServer = prepareAsyncServerBuilder().serverInfo("test-server", "1.0.0").tools(tool).build();
-		var mcpServer = prepareAsyncServerBuilder()
-				.serverInfo("test-server", "1.0.0")
-				.tools(tool)
-				.build();
 
-		try (
+		try (var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0"))
-			var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0"))
+			.capabilities(ClientCapabilities.builder().sampling().build())
-				.capabilities(ClientCapabilities.builder().sampling().build())
+			.sampling(samplingHandler)
-				.sampling(samplingHandler)
+			.build()) {
-				.build()) {//@formatter:on
 
 			InitializeResult initResult = mcpClient.initialize();
 			assertThat(initResult).isNotNull();
 
 			CallToolResult response = mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of()));
 
-			assertThat(response).isNotNull().isEqualTo(callResponse);
+			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();
 	}
@@ -225,6 +226,8 @@ public abstract class AbstractMcpClientServerIntegrationTests {
 		CallToolResult callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")),
 				null);
 
+		AtomicReference<CreateMessageResult> samplingResult = new AtomicReference<>();
+
 		McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder()
 			.tool(Tool.builder().name("tool1").description("tool1 description").inputSchema(emptyJsonSchema).build())
 			.callHandler((exchange, request) -> {
@@ -240,16 +243,9 @@ public abstract class AbstractMcpClientServerIntegrationTests {
 						.build())
 					.build();
 
-				StepVerifier.create(exchange.createMessage(createMessageRequest)).consumeNextWith(result -> {
+				return exchange.createMessage(createMessageRequest)
-					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);
 			})
 			.build();
 
@@ -266,6 +262,15 @@ public abstract class AbstractMcpClientServerIntegrationTests {
 		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);
+		});
+
 		mcpClient.close();
 		mcpServer.close();
 	}
@@ -312,16 +317,7 @@ public abstract class AbstractMcpClientServerIntegrationTests {
 						.build())
 					.build();
 
-				StepVerifier.create(exchange.createMessage(createMessageRequest)).consumeNextWith(result -> {
+				return exchange.createMessage(createMessageRequest).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);
 			})
 			.build();
 
@@ -335,7 +331,7 @@ public abstract class AbstractMcpClientServerIntegrationTests {
 
 		assertThatExceptionOfType(McpError.class).isThrownBy(() -> {
 			mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of()));
-		}).withMessageContaining("Timeout");
+		}).withMessageContaining("1000ms");
 
 		mcpClient.close();
 		mcpServer.close();
@@ -352,19 +348,14 @@ public abstract class AbstractMcpClientServerIntegrationTests {
 
 		McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder()
 			.tool(Tool.builder().name("tool1").description("tool1 description").inputSchema(emptyJsonSchema).build())
-			.callHandler((exchange, request) -> {
+			.callHandler((exchange, request) -> exchange.createElicitation(mock(ElicitRequest.class))
-
+				.then(Mono.just(mock(CallToolResult.class))))
-				exchange.createElicitation(mock(McpSchema.ElicitRequest.class)).block();
-
-				return Mono.just(mock(CallToolResult.class));
-			})
 			.build();
 
 		var server = prepareAsyncServerBuilder().serverInfo("test-server", "1.0.0").tools(tool).build();
 
-		try (
+		// Create client without elicitation capabilities
-				// Create client without elicitation capabilities
+		try (var client = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0")).build()) {
-				var client = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0")).build()) {
 
 			assertThat(client.initialize()).isNotNull();
 
@@ -440,17 +431,10 @@ public abstract class AbstractMcpClientServerIntegrationTests {
 
 		var clientBuilder = clientBuilders.get(clientType);
 
-		Function<McpSchema.ElicitRequest, McpSchema.ElicitResult> elicitationHandler = request -> {
+		Function<ElicitRequest, ElicitResult> elicitationHandler = request -> {
 			assertThat(request.message()).isNotEmpty();
 			assertThat(request.requestedSchema()).isNotNull();
-			try {
+			return new ElicitResult(ElicitResult.Action.ACCEPT, Map.of("message", request.message()));
-				TimeUnit.SECONDS.sleep(2);
-			}
-			catch (InterruptedException e) {
-				throw new RuntimeException(e);
-			}
-			return new McpSchema.ElicitResult(McpSchema.ElicitResult.Action.ACCEPT,
-					Map.of("message", request.message()));
 		};
 
 		var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0"))
@@ -461,6 +445,8 @@ public abstract class AbstractMcpClientServerIntegrationTests {
 		CallToolResult callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")),
 				null);
 
+		AtomicReference<ElicitResult> resultRef = new AtomicReference<>();
+
 		McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder()
 			.tool(Tool.builder().name("tool1").description("tool1 description").inputSchema(emptyJsonSchema).build())
 			.callHandler((exchange, request) -> {
@@ -471,13 +457,9 @@ public abstract class AbstractMcpClientServerIntegrationTests {
 							Map.of("type", "object", "properties", Map.of("message", Map.of("type", "string"))))
 					.build();
 
-				StepVerifier.create(exchange.createElicitation(elicitationRequest)).consumeNextWith(result -> {
+				return exchange.createElicitation(elicitationRequest)
-					assertThat(result).isNotNull();
+					.doOnNext(resultRef::set)
-					assertThat(result.action()).isEqualTo(McpSchema.ElicitResult.Action.ACCEPT);
+					.then(Mono.just(callResponse));
-					assertThat(result.content().get("message")).isEqualTo("Test message");
-				}).verifyComplete();
-
-				return Mono.just(callResponse);
 			})
 			.build();
 
@@ -493,6 +475,11 @@ public abstract class AbstractMcpClientServerIntegrationTests {
 
 		assertThat(response).isNotNull();
 		assertThat(response).isEqualTo(callResponse);
+		assertWith(resultRef.get(), result -> {
+			assertThat(result).isNotNull();
+			assertThat(result.action()).isEqualTo(McpSchema.ElicitResult.Action.ACCEPT);
+			assertThat(result.content().get("message")).isEqualTo("Test message");
+		});
 
 		mcpClient.closeGracefully();
 		mcpServer.closeGracefully().block();
@@ -870,7 +857,7 @@ public abstract class AbstractMcpClientServerIntegrationTests {
 			})
 			.build();
 
-		AtomicReference<List<Tool>> rootsRef = new AtomicReference<>();
+		AtomicReference<List<Tool>> toolsRef = new AtomicReference<>();
 
 		var mcpServer = prepareSyncServerBuilder().capabilities(ServerCapabilities.builder().tools(true).build())
 			.tools(tool1)
@@ -887,32 +874,31 @@ public abstract class AbstractMcpClientServerIntegrationTests {
 						.build(), HttpResponse.BodyHandlers.ofString());
 				String responseBody = response.body();
 				assertThat(responseBody).isNotBlank();
+				toolsRef.set(toolsUpdate);
 			}
 			catch (Exception e) {
 				e.printStackTrace();
 			}
-
-			rootsRef.set(toolsUpdate);
 		}).build()) {
 
 			InitializeResult initResult = mcpClient.initialize();
 			assertThat(initResult).isNotNull();
 
-			assertThat(rootsRef.get()).isNull();
+			assertThat(toolsRef.get()).isNull();
 
 			assertThat(mcpClient.listTools().tools()).contains(tool1.tool());
 
 			mcpServer.notifyToolsListChanged();
 
 			await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> {
-				assertThat(rootsRef.get()).containsAll(List.of(tool1.tool()));
+				assertThat(toolsRef.get()).containsAll(List.of(tool1.tool()));
 			});
 
 			// Remove a tool
 			mcpServer.removeTool("tool1");
 
 			await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> {
-				assertThat(rootsRef.get()).isEmpty();
+				assertThat(toolsRef.get()).isEmpty();
 			});
 
 			// Add a new tool
@@ -928,7 +914,7 @@ public abstract class AbstractMcpClientServerIntegrationTests {
 			mcpServer.addTool(tool2);
 
 			await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> {
-				assertThat(rootsRef.get()).containsAll(List.of(tool2.tool()));
+				assertThat(toolsRef.get()).containsAll(List.of(tool2.tool()));
 			});
 		}
 

mcp/src/test/java/io/modelcontextprotocol/server/AbstractMcpClientServerIntegrationTests.java

@@ -8,6 +8,7 @@ import static net.javacrumbs.jsonunit.assertj.JsonAssertions.assertThatJson;
 import static net.javacrumbs.jsonunit.assertj.JsonAssertions.json;
 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;
 
@@ -43,6 +44,9 @@ import io.modelcontextprotocol.spec.McpSchema.ElicitRequest;
 import io.modelcontextprotocol.spec.McpSchema.ElicitResult;
 import io.modelcontextprotocol.spec.McpSchema.InitializeResult;
 import io.modelcontextprotocol.spec.McpSchema.ModelPreferences;
+import io.modelcontextprotocol.spec.McpSchema.Prompt;
+import io.modelcontextprotocol.spec.McpSchema.PromptArgument;
+import io.modelcontextprotocol.spec.McpSchema.PromptReference;
 import io.modelcontextprotocol.spec.McpSchema.Role;
 import io.modelcontextprotocol.spec.McpSchema.Root;
 import io.modelcontextprotocol.spec.McpSchema.ServerCapabilities;
@@ -137,6 +141,8 @@ public abstract class AbstractMcpClientServerIntegrationTests {
 		CallToolResult callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")),
 				null);
 
+		AtomicReference<CreateMessageResult> samplingResult = new AtomicReference<>();
+
 		McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder()
 			.tool(Tool.builder().name("tool1").description("tool1 description").inputSchema(emptyJsonSchema).build())
 			.callHandler((exchange, request) -> {
@@ -152,37 +158,35 @@ public abstract class AbstractMcpClientServerIntegrationTests {
 						.build())
 					.build();
 
-				StepVerifier.create(exchange.createMessage(createMessageRequest)).consumeNextWith(result -> {
+				return exchange.createMessage(createMessageRequest)
-					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);
 			})
 			.build();
 
-		//@formatter:off		
+		var mcpServer = prepareAsyncServerBuilder().serverInfo("test-server", "1.0.0").tools(tool).build();
-		var mcpServer = prepareAsyncServerBuilder()
-				.serverInfo("test-server", "1.0.0")
-				.tools(tool)
-				.build();
 
-		try (
+		try (var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0"))
-			var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0"))
+			.capabilities(ClientCapabilities.builder().sampling().build())
-				.capabilities(ClientCapabilities.builder().sampling().build())
+			.sampling(samplingHandler)
-				.sampling(samplingHandler)
+			.build()) {
-				.build()) {//@formatter:on
 
 			InitializeResult initResult = mcpClient.initialize();
 			assertThat(initResult).isNotNull();
 
 			CallToolResult response = mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of()));
 
-			assertThat(response).isNotNull().isEqualTo(callResponse);
+			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();
 	}
@@ -218,6 +222,8 @@ public abstract class AbstractMcpClientServerIntegrationTests {
 		CallToolResult callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")),
 				null);
 
+		AtomicReference<CreateMessageResult> samplingResult = new AtomicReference<>();
+
 		McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder()
 			.tool(Tool.builder().name("tool1").description("tool1 description").inputSchema(emptyJsonSchema).build())
 			.callHandler((exchange, request) -> {
@@ -233,16 +239,9 @@ public abstract class AbstractMcpClientServerIntegrationTests {
 						.build())
 					.build();
 
-				StepVerifier.create(exchange.createMessage(createMessageRequest)).consumeNextWith(result -> {
+				return exchange.createMessage(createMessageRequest)
-					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);
 			})
 			.build();
 
@@ -259,6 +258,15 @@ public abstract class AbstractMcpClientServerIntegrationTests {
 		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);
+		});
+
 		mcpClient.close();
 		mcpServer.close();
 	}
@@ -305,16 +313,7 @@ public abstract class AbstractMcpClientServerIntegrationTests {
 						.build())
 					.build();
 
-				StepVerifier.create(exchange.createMessage(createMessageRequest)).consumeNextWith(result -> {
+				return exchange.createMessage(createMessageRequest).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);
 			})
 			.build();
 
@@ -328,7 +327,7 @@ public abstract class AbstractMcpClientServerIntegrationTests {
 
 		assertThatExceptionOfType(McpError.class).isThrownBy(() -> {
 			mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of()));
-		}).withMessageContaining("Timeout");
+		}).withMessageContaining("1000ms");
 
 		mcpClient.close();
 		mcpServer.close();
@@ -345,19 +344,14 @@ public abstract class AbstractMcpClientServerIntegrationTests {
 
 		McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder()
 			.tool(Tool.builder().name("tool1").description("tool1 description").inputSchema(emptyJsonSchema).build())
-			.callHandler((exchange, request) -> {
+			.callHandler((exchange, request) -> exchange.createElicitation(mock(ElicitRequest.class))
-
+				.then(Mono.just(mock(CallToolResult.class))))
-				exchange.createElicitation(mock(McpSchema.ElicitRequest.class)).block();
-
-				return Mono.just(mock(CallToolResult.class));
-			})
 			.build();
 
 		var server = prepareAsyncServerBuilder().serverInfo("test-server", "1.0.0").tools(tool).build();
 
-		try (
+		// Create client without elicitation capabilities
-				// Create client without elicitation capabilities
+		try (var client = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0")).build()) {
-				var client = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0")).build()) {
 
 			assertThat(client.initialize()).isNotNull();
 
@@ -433,17 +427,10 @@ public abstract class AbstractMcpClientServerIntegrationTests {
 
 		var clientBuilder = clientBuilders.get(clientType);
 
-		Function<McpSchema.ElicitRequest, McpSchema.ElicitResult> elicitationHandler = request -> {
+		Function<ElicitRequest, ElicitResult> elicitationHandler = request -> {
 			assertThat(request.message()).isNotEmpty();
 			assertThat(request.requestedSchema()).isNotNull();
-			try {
+			return new ElicitResult(ElicitResult.Action.ACCEPT, Map.of("message", request.message()));
-				TimeUnit.SECONDS.sleep(2);
-			}
-			catch (InterruptedException e) {
-				throw new RuntimeException(e);
-			}
-			return new McpSchema.ElicitResult(McpSchema.ElicitResult.Action.ACCEPT,
-					Map.of("message", request.message()));
 		};
 
 		var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0"))
@@ -454,6 +441,8 @@ public abstract class AbstractMcpClientServerIntegrationTests {
 		CallToolResult callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")),
 				null);
 
+		AtomicReference<ElicitResult> resultRef = new AtomicReference<>();
+
 		McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder()
 			.tool(Tool.builder().name("tool1").description("tool1 description").inputSchema(emptyJsonSchema).build())
 			.callHandler((exchange, request) -> {
@@ -464,13 +453,9 @@ public abstract class AbstractMcpClientServerIntegrationTests {
 							Map.of("type", "object", "properties", Map.of("message", Map.of("type", "string"))))
 					.build();
 
-				StepVerifier.create(exchange.createElicitation(elicitationRequest)).consumeNextWith(result -> {
+				return exchange.createElicitation(elicitationRequest)
-					assertThat(result).isNotNull();
+					.doOnNext(resultRef::set)
-					assertThat(result.action()).isEqualTo(McpSchema.ElicitResult.Action.ACCEPT);
+					.then(Mono.just(callResponse));
-					assertThat(result.content().get("message")).isEqualTo("Test message");
-				}).verifyComplete();
-
-				return Mono.just(callResponse);
 			})
 			.build();
 
@@ -486,6 +471,11 @@ public abstract class AbstractMcpClientServerIntegrationTests {
 
 		assertThat(response).isNotNull();
 		assertThat(response).isEqualTo(callResponse);
+		assertWith(resultRef.get(), result -> {
+			assertThat(result).isNotNull();
+			assertThat(result.action()).isEqualTo(McpSchema.ElicitResult.Action.ACCEPT);
+			assertThat(result.content().get("message")).isEqualTo("Test message");
+		});
 
 		mcpClient.closeGracefully();
 		mcpServer.closeGracefully().block();
@@ -863,7 +853,7 @@ public abstract class AbstractMcpClientServerIntegrationTests {
 			})
 			.build();
 
-		AtomicReference<List<Tool>> rootsRef = new AtomicReference<>();
+		AtomicReference<List<Tool>> toolsRef = new AtomicReference<>();
 
 		var mcpServer = prepareSyncServerBuilder().capabilities(ServerCapabilities.builder().tools(true).build())
 			.tools(tool1)
@@ -880,32 +870,31 @@ public abstract class AbstractMcpClientServerIntegrationTests {
 						.build(), HttpResponse.BodyHandlers.ofString());
 				String responseBody = response.body();
 				assertThat(responseBody).isNotBlank();
+				toolsRef.set(toolsUpdate);
 			}
 			catch (Exception e) {
 				e.printStackTrace();
 			}
-
-			rootsRef.set(toolsUpdate);
 		}).build()) {
 
 			InitializeResult initResult = mcpClient.initialize();
 			assertThat(initResult).isNotNull();
 
-			assertThat(rootsRef.get()).isNull();
+			assertThat(toolsRef.get()).isNull();
 
 			assertThat(mcpClient.listTools().tools()).contains(tool1.tool());
 
 			mcpServer.notifyToolsListChanged();
 
 			await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> {
-				assertThat(rootsRef.get()).containsAll(List.of(tool1.tool()));
+				assertThat(toolsRef.get()).containsAll(List.of(tool1.tool()));
 			});
 
 			// Remove a tool
 			mcpServer.removeTool("tool1");
 
 			await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> {
-				assertThat(rootsRef.get()).isEmpty();
+				assertThat(toolsRef.get()).isEmpty();
 			});
 
 			// Add a new tool
@@ -921,7 +910,7 @@ public abstract class AbstractMcpClientServerIntegrationTests {
 			mcpServer.addTool(tool2);
 
 			await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> {
-				assertThat(rootsRef.get()).containsAll(List.of(tool2.tool()));
+				assertThat(toolsRef.get()).containsAll(List.of(tool2.tool()));
 			});
 		}
 
@@ -1184,8 +1173,8 @@ public abstract class AbstractMcpClientServerIntegrationTests {
 
 		var mcpServer = prepareSyncServerBuilder().capabilities(ServerCapabilities.builder().completions().build())
 			.prompts(new McpServerFeatures.SyncPromptSpecification(
-					new McpSchema.Prompt("code_review", "Code review", "this is code review prompt",
+					new Prompt("code_review", "Code review", "this is code review prompt",
-							List.of(new McpSchema.PromptArgument("language", "Language", "string", false))),
+							List.of(new PromptArgument("language", "Language", "string", false))),
 					(mcpSyncServerExchange, getPromptRequest) -> null))
 			.completions(new McpServerFeatures.SyncCompletionSpecification(
 					new McpSchema.PromptReference("ref/prompt", "code_review", "Code review"), completionHandler))
@@ -1197,7 +1186,7 @@ public abstract class AbstractMcpClientServerIntegrationTests {
 			assertThat(initResult).isNotNull();
 
 			CompleteRequest request = new CompleteRequest(
-					new McpSchema.PromptReference("ref/prompt", "code_review", "Code review"),
+					new PromptReference("ref/prompt", "code_review", "Code review"),
 					new CompleteRequest.CompleteArgument("language", "py"));
 
 			CompleteResult result = mcpClient.completeCompletion(request);