Skip to content

Commit ed2c35e

Browse files
tzolovtzolov
committed
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 <[email protected]>
1 parent dee8411 commit ed2c35e

File tree

3 files changed

+126
-154
lines changed

3 files changed

+126
-154
lines changed

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

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -11,8 +11,6 @@
1111
import org.apache.catalina.LifecycleState;
1212
import org.junit.jupiter.api.AfterEach;
1313
import org.junit.jupiter.api.BeforeEach;
14-
import org.junit.jupiter.params.ParameterizedTest;
15-
import org.junit.jupiter.params.provider.ValueSource;
1614
import org.springframework.context.annotation.Bean;
1715
import org.springframework.context.annotation.Configuration;
1816
import org.springframework.web.reactive.function.client.WebClient;
@@ -29,7 +27,6 @@
2927
import io.modelcontextprotocol.server.McpServer.AsyncSpecification;
3028
import io.modelcontextprotocol.server.McpServer.SyncSpecification;
3129
import io.modelcontextprotocol.server.transport.WebMvcStreamableServerTransportProvider;
32-
import io.modelcontextprotocol.spec.McpSchema;
3330
import reactor.core.scheduler.Schedulers;
3431

3532
class WebMvcStreamableIntegrationTests extends AbstractMcpClientServerIntegrationTests {

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

Lines changed: 60 additions & 74 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
import static net.javacrumbs.jsonunit.assertj.JsonAssertions.json;
99
import static org.assertj.core.api.Assertions.assertThat;
1010
import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
11+
import static org.assertj.core.api.Assertions.assertWith;
1112
import static org.awaitility.Awaitility.await;
1213
import static org.mockito.Mockito.mock;
1314

@@ -144,6 +145,8 @@ void testCreateMessageSuccess(String clientType) {
144145
CallToolResult callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")),
145146
null);
146147

148+
AtomicReference<CreateMessageResult> samplingResult = new AtomicReference<>();
149+
147150
McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder()
148151
.tool(Tool.builder().name("tool1").description("tool1 description").inputSchema(emptyJsonSchema).build())
149152
.callHandler((exchange, request) -> {
@@ -159,37 +162,35 @@ void testCreateMessageSuccess(String clientType) {
159162
.build())
160163
.build();
161164

162-
StepVerifier.create(exchange.createMessage(createMessageRequest)).consumeNextWith(result -> {
163-
assertThat(result).isNotNull();
164-
assertThat(result.role()).isEqualTo(Role.USER);
165-
assertThat(result.content()).isInstanceOf(McpSchema.TextContent.class);
166-
assertThat(((McpSchema.TextContent) result.content()).text()).isEqualTo("Test message");
167-
assertThat(result.model()).isEqualTo("MockModelName");
168-
assertThat(result.stopReason()).isEqualTo(CreateMessageResult.StopReason.STOP_SEQUENCE);
169-
}).verifyComplete();
170-
171-
return Mono.just(callResponse);
165+
return exchange.createMessage(createMessageRequest)
166+
.doOnNext(samplingResult::set)
167+
.thenReturn(callResponse);
172168
})
173169
.build();
174170

175-
//@formatter:off
176-
var mcpServer = prepareAsyncServerBuilder()
177-
.serverInfo("test-server", "1.0.0")
178-
.tools(tool)
179-
.build();
171+
var mcpServer = prepareAsyncServerBuilder().serverInfo("test-server", "1.0.0").tools(tool).build();
180172

181-
try (
182-
var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0"))
183-
.capabilities(ClientCapabilities.builder().sampling().build())
184-
.sampling(samplingHandler)
185-
.build()) {//@formatter:on
173+
try (var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0"))
174+
.capabilities(ClientCapabilities.builder().sampling().build())
175+
.sampling(samplingHandler)
176+
.build()) {
186177

187178
InitializeResult initResult = mcpClient.initialize();
188179
assertThat(initResult).isNotNull();
189180

190181
CallToolResult response = mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of()));
191182

192-
assertThat(response).isNotNull().isEqualTo(callResponse);
183+
assertThat(response).isNotNull();
184+
assertThat(response).isEqualTo(callResponse);
185+
186+
assertWith(samplingResult.get(), result -> {
187+
assertThat(result).isNotNull();
188+
assertThat(result.role()).isEqualTo(Role.USER);
189+
assertThat(result.content()).isInstanceOf(McpSchema.TextContent.class);
190+
assertThat(((McpSchema.TextContent) result.content()).text()).isEqualTo("Test message");
191+
assertThat(result.model()).isEqualTo("MockModelName");
192+
assertThat(result.stopReason()).isEqualTo(CreateMessageResult.StopReason.STOP_SEQUENCE);
193+
});
193194
}
194195
mcpServer.close();
195196
}
@@ -225,6 +226,8 @@ void testCreateMessageWithRequestTimeoutSuccess(String clientType) throws Interr
225226
CallToolResult callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")),
226227
null);
227228

229+
AtomicReference<CreateMessageResult> samplingResult = new AtomicReference<>();
230+
228231
McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder()
229232
.tool(Tool.builder().name("tool1").description("tool1 description").inputSchema(emptyJsonSchema).build())
230233
.callHandler((exchange, request) -> {
@@ -240,16 +243,9 @@ void testCreateMessageWithRequestTimeoutSuccess(String clientType) throws Interr
240243
.build())
241244
.build();
242245

243-
StepVerifier.create(exchange.createMessage(createMessageRequest)).consumeNextWith(result -> {
244-
assertThat(result).isNotNull();
245-
assertThat(result.role()).isEqualTo(Role.USER);
246-
assertThat(result.content()).isInstanceOf(McpSchema.TextContent.class);
247-
assertThat(((McpSchema.TextContent) result.content()).text()).isEqualTo("Test message");
248-
assertThat(result.model()).isEqualTo("MockModelName");
249-
assertThat(result.stopReason()).isEqualTo(CreateMessageResult.StopReason.STOP_SEQUENCE);
250-
}).verifyComplete();
251-
252-
return Mono.just(callResponse);
246+
return exchange.createMessage(createMessageRequest)
247+
.doOnNext(samplingResult::set)
248+
.thenReturn(callResponse);
253249
})
254250
.build();
255251

@@ -266,6 +262,15 @@ void testCreateMessageWithRequestTimeoutSuccess(String clientType) throws Interr
266262
assertThat(response).isNotNull();
267263
assertThat(response).isEqualTo(callResponse);
268264

265+
assertWith(samplingResult.get(), result -> {
266+
assertThat(result).isNotNull();
267+
assertThat(result.role()).isEqualTo(Role.USER);
268+
assertThat(result.content()).isInstanceOf(McpSchema.TextContent.class);
269+
assertThat(((McpSchema.TextContent) result.content()).text()).isEqualTo("Test message");
270+
assertThat(result.model()).isEqualTo("MockModelName");
271+
assertThat(result.stopReason()).isEqualTo(CreateMessageResult.StopReason.STOP_SEQUENCE);
272+
});
273+
269274
mcpClient.close();
270275
mcpServer.close();
271276
}
@@ -312,16 +317,7 @@ void testCreateMessageWithRequestTimeoutFail(String clientType) throws Interrupt
312317
.build())
313318
.build();
314319

315-
StepVerifier.create(exchange.createMessage(createMessageRequest)).consumeNextWith(result -> {
316-
assertThat(result).isNotNull();
317-
assertThat(result.role()).isEqualTo(Role.USER);
318-
assertThat(result.content()).isInstanceOf(McpSchema.TextContent.class);
319-
assertThat(((McpSchema.TextContent) result.content()).text()).isEqualTo("Test message");
320-
assertThat(result.model()).isEqualTo("MockModelName");
321-
assertThat(result.stopReason()).isEqualTo(CreateMessageResult.StopReason.STOP_SEQUENCE);
322-
}).verifyComplete();
323-
324-
return Mono.just(callResponse);
320+
return exchange.createMessage(createMessageRequest).thenReturn(callResponse);
325321
})
326322
.build();
327323

@@ -335,7 +331,7 @@ void testCreateMessageWithRequestTimeoutFail(String clientType) throws Interrupt
335331

336332
assertThatExceptionOfType(McpError.class).isThrownBy(() -> {
337333
mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of()));
338-
}).withMessageContaining("Timeout");
334+
}).withMessageContaining("1000ms");
339335

340336
mcpClient.close();
341337
mcpServer.close();
@@ -352,19 +348,14 @@ void testCreateElicitationWithoutElicitationCapabilities(String clientType) {
352348

353349
McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder()
354350
.tool(Tool.builder().name("tool1").description("tool1 description").inputSchema(emptyJsonSchema).build())
355-
.callHandler((exchange, request) -> {
356-
357-
exchange.createElicitation(mock(McpSchema.ElicitRequest.class)).block();
358-
359-
return Mono.just(mock(CallToolResult.class));
360-
})
351+
.callHandler((exchange, request) -> exchange.createElicitation(mock(ElicitRequest.class))
352+
.then(Mono.just(mock(CallToolResult.class))))
361353
.build();
362354

363355
var server = prepareAsyncServerBuilder().serverInfo("test-server", "1.0.0").tools(tool).build();
364356

365-
try (
366-
// Create client without elicitation capabilities
367-
var client = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0")).build()) {
357+
// Create client without elicitation capabilities
358+
try (var client = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0")).build()) {
368359

369360
assertThat(client.initialize()).isNotNull();
370361

@@ -440,17 +431,10 @@ void testCreateElicitationWithRequestTimeoutSuccess(String clientType) {
440431

441432
var clientBuilder = clientBuilders.get(clientType);
442433

443-
Function<McpSchema.ElicitRequest, McpSchema.ElicitResult> elicitationHandler = request -> {
434+
Function<ElicitRequest, ElicitResult> elicitationHandler = request -> {
444435
assertThat(request.message()).isNotEmpty();
445436
assertThat(request.requestedSchema()).isNotNull();
446-
try {
447-
TimeUnit.SECONDS.sleep(2);
448-
}
449-
catch (InterruptedException e) {
450-
throw new RuntimeException(e);
451-
}
452-
return new McpSchema.ElicitResult(McpSchema.ElicitResult.Action.ACCEPT,
453-
Map.of("message", request.message()));
437+
return new ElicitResult(ElicitResult.Action.ACCEPT, Map.of("message", request.message()));
454438
};
455439

456440
var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0"))
@@ -461,6 +445,8 @@ void testCreateElicitationWithRequestTimeoutSuccess(String clientType) {
461445
CallToolResult callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")),
462446
null);
463447

448+
AtomicReference<ElicitResult> resultRef = new AtomicReference<>();
449+
464450
McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder()
465451
.tool(Tool.builder().name("tool1").description("tool1 description").inputSchema(emptyJsonSchema).build())
466452
.callHandler((exchange, request) -> {
@@ -471,13 +457,9 @@ void testCreateElicitationWithRequestTimeoutSuccess(String clientType) {
471457
Map.of("type", "object", "properties", Map.of("message", Map.of("type", "string"))))
472458
.build();
473459

474-
StepVerifier.create(exchange.createElicitation(elicitationRequest)).consumeNextWith(result -> {
475-
assertThat(result).isNotNull();
476-
assertThat(result.action()).isEqualTo(McpSchema.ElicitResult.Action.ACCEPT);
477-
assertThat(result.content().get("message")).isEqualTo("Test message");
478-
}).verifyComplete();
479-
480-
return Mono.just(callResponse);
460+
return exchange.createElicitation(elicitationRequest)
461+
.doOnNext(resultRef::set)
462+
.then(Mono.just(callResponse));
481463
})
482464
.build();
483465

@@ -493,6 +475,11 @@ void testCreateElicitationWithRequestTimeoutSuccess(String clientType) {
493475

494476
assertThat(response).isNotNull();
495477
assertThat(response).isEqualTo(callResponse);
478+
assertWith(resultRef.get(), result -> {
479+
assertThat(result).isNotNull();
480+
assertThat(result.action()).isEqualTo(McpSchema.ElicitResult.Action.ACCEPT);
481+
assertThat(result.content().get("message")).isEqualTo("Test message");
482+
});
496483

497484
mcpClient.closeGracefully();
498485
mcpServer.closeGracefully().block();
@@ -870,7 +857,7 @@ void testToolListChangeHandlingSuccess(String clientType) {
870857
})
871858
.build();
872859

873-
AtomicReference<List<Tool>> rootsRef = new AtomicReference<>();
860+
AtomicReference<List<Tool>> toolsRef = new AtomicReference<>();
874861

875862
var mcpServer = prepareSyncServerBuilder().capabilities(ServerCapabilities.builder().tools(true).build())
876863
.tools(tool1)
@@ -887,32 +874,31 @@ void testToolListChangeHandlingSuccess(String clientType) {
887874
.build(), HttpResponse.BodyHandlers.ofString());
888875
String responseBody = response.body();
889876
assertThat(responseBody).isNotBlank();
877+
toolsRef.set(toolsUpdate);
890878
}
891879
catch (Exception e) {
892880
e.printStackTrace();
893881
}
894-
895-
rootsRef.set(toolsUpdate);
896882
}).build()) {
897883

898884
InitializeResult initResult = mcpClient.initialize();
899885
assertThat(initResult).isNotNull();
900886

901-
assertThat(rootsRef.get()).isNull();
887+
assertThat(toolsRef.get()).isNull();
902888

903889
assertThat(mcpClient.listTools().tools()).contains(tool1.tool());
904890

905891
mcpServer.notifyToolsListChanged();
906892

907893
await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> {
908-
assertThat(rootsRef.get()).containsAll(List.of(tool1.tool()));
894+
assertThat(toolsRef.get()).containsAll(List.of(tool1.tool()));
909895
});
910896

911897
// Remove a tool
912898
mcpServer.removeTool("tool1");
913899

914900
await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> {
915-
assertThat(rootsRef.get()).isEmpty();
901+
assertThat(toolsRef.get()).isEmpty();
916902
});
917903

918904
// Add a new tool
@@ -928,7 +914,7 @@ void testToolListChangeHandlingSuccess(String clientType) {
928914
mcpServer.addTool(tool2);
929915

930916
await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> {
931-
assertThat(rootsRef.get()).containsAll(List.of(tool2.tool()));
917+
assertThat(toolsRef.get()).containsAll(List.of(tool2.tool()));
932918
});
933919
}
934920

0 commit comments

Comments
 (0)