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

refactor: improve integration tests stability

Browse source 
- 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>
This commit is contained in:
Christian Tzolov 2025-08-09 18:41:01 +01:00
commit eb427add79
1 changed files with 126 additions and 154 deletions
Download patch file Download diff file Expand all files Collapse all files

143
mcp/src/test/java/io/modelcontextprotocol/server/AbstractMcpClientServerIntegrationTests.java
View file

@ -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 -> {
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(createMessageRequest)
.doOnNext(samplingResult::set)
.thenReturn(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 (
var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0"))
.capabilities(ClientCapabilities.builder().sampling().build())
.sampling(samplingHandler)
.build()) {//@formatter:on
try (var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0"))
.capabilities(ClientCapabilities.builder().sampling().build())
.sampling(samplingHandler)
.build()) {
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 -> {
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(createMessageRequest)
.doOnNext(samplingResult::set)
.thenReturn(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 -> {
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(createMessageRequest).thenReturn(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) -> {
exchange.createElicitation(mock(McpSchema.ElicitRequest.class)).block();
return Mono.just(mock(CallToolResult.class));
})
.callHandler((exchange, request) -> exchange.createElicitation(mock(ElicitRequest.class))
.then(Mono.just(mock(CallToolResult.class))))
.build();
var server = prepareAsyncServerBuilder().serverInfo("test-server", "1.0.0").tools(tool).build();
try (
// Create client without elicitation capabilities
var client = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0")).build()) {
// Create client without elicitation capabilities
try (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 {
TimeUnit.SECONDS.sleep(2);
}
catch (InterruptedException e) {
throw new RuntimeException(e);
}
return new McpSchema.ElicitResult(McpSchema.ElicitResult.Action.ACCEPT,
Map.of("message", request.message()));
return new ElicitResult(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 -> {
assertThat(result).isNotNull();
assertThat(result.action()).isEqualTo(McpSchema.ElicitResult.Action.ACCEPT);
assertThat(result.content().get("message")).isEqualTo("Test message");
}).verifyComplete();
return Mono.just(callResponse);
return exchange.createElicitation(elicitationRequest)
.doOnNext(resultRef::set)
.then(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",
List.of(new McpSchema.PromptArgument("language", "Language", "string", false))),
new Prompt("code_review", "Code review", "this is code review prompt",
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);