modelcontextprotocol/java-sdk
Watch
1
0
Fork
You've already forked java-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

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

Browse source 
- 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>
This commit is contained in:
Dariusz Jędrzejczyk 2025-04-14 12:08:30 +02:00 • committed by Christian Tzolov
commit f7f8ccd0ac
1 changed files with 62 additions and 73 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -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) -> {
exchange.createMessage(mock(McpSchema.CreateMessageRequest.class)).block();
return Mono.just(mock(CallToolResult.class));
});
new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema),
(exchange, request) -> exchange.createMessage(mock(CreateMessageRequest.class))
.thenReturn(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 -> {
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);
return exchange.createMessage(craeteMessageRequest)
.doOnNext(samplingResult::set)
.thenReturn(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 -> {
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);
return exchange.createMessage(craeteMessageRequest)
.doOnNext(samplingResult::set)
.thenReturn(callResponse);
});
var mcpServer = McpServer.async(mcpServerTransportProvider)
@ -260,16 +250,30 @@ class WebFluxSseIntegrationTests {
.tools(tool)
.build();
InitializeResult initResult = mcpClient.initialize();
assertThat(initResult).isNotNull();
try (var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0"))
.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();
assertThat(response).isEqualTo(callResponse);
CallToolResult response = mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of()));
mcpClient.close();
mcpServer.close();
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.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 -> {
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);
return exchange.createMessage(craeteMessageRequest).thenReturn(callResponse);
});
var mcpServer = McpServer.async(mcpServerTransportProvider)
@ -334,15 +318,21 @@ class WebFluxSseIntegrationTests {
.tools(tool)
.build();
InitializeResult initResult = mcpClient.initialize();
assertThat(initResult).isNotNull();
try (var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0"))
.capabilities(ClientCapabilities.builder().sampling().build())
.sampling(samplingHandler)
.build()) {
assertThatExceptionOfType(McpError.class).isThrownBy(() -> {
mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of()));
}).withMessageContaining("Timeout");
InitializeResult initResult = mcpClient.initialize();
assertThat(initResult).isNotNull();
mcpClient.close();
mcpServer.close();
assertThatExceptionOfType(McpError.class).isThrownBy(() -> {
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
var mcpClient = clientBuilder.capabilities(ClientCapabilities.builder().build()).build()) {
// Create client without roots capability
try (var mcpClient = clientBuilder.capabilities(ClientCapabilities.builder().build()).build()) {
assertThat(mcpClient.initialize()).isNotNull();