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 WebFluxSse integration test

Browse source 
Signed-off-by: Dariusz Jędrzejczyk <dariusz.jedrzejczyk@broadcom.com>
This commit is contained in:
Dariusz Jędrzejczyk 2025-05-14 22:46:54 +02:00 • committed by Christian Tzolov
commit 2e13f9f9df
1 changed files with 24 additions and 22 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -8,6 +8,8 @@ import java.util.ArrayList;
import java.util.List; import java.util.List;
import java.util.Map; import java.util.Map;
import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference; import java.util.concurrent.atomic.AtomicReference;
import java.util.function.BiFunction; import java.util.function.BiFunction;
@ -651,9 +653,11 @@ class WebFluxSseIntegrationTests {
@ParameterizedTest(name = "{0} : {displayName} ") @ParameterizedTest(name = "{0} : {displayName} ")
@ValueSource(strings = { "httpclient", "webflux" }) @ValueSource(strings = { "httpclient", "webflux" })
void testLoggingNotification(String clientType) { void testLoggingNotification(String clientType) throws InterruptedException {
int expectedNotificationsCount = 3;
CountDownLatch latch = new CountDownLatch(expectedNotificationsCount);
// Create a list to store received logging notifications // Create a list to store received logging notifications
List<McpSchema.LoggingMessageNotification> receivedNotifications = new ArrayList<>(); List<McpSchema.LoggingMessageNotification> receivedNotifications = new CopyOnWriteArrayList<>();
var clientBuilder = clientBuilders.get(clientType); var clientBuilder = clientBuilders.get(clientType);
@ -709,6 +713,7 @@ class WebFluxSseIntegrationTests {
// Create client with logging notification handler // Create client with logging notification handler
var mcpClient = clientBuilder.loggingConsumer(notification -> { var mcpClient = clientBuilder.loggingConsumer(notification -> {
receivedNotifications.add(notification); receivedNotifications.add(notification);
latch.countDown();
}).build()) { }).build()) {
// Initialize client // Initialize client
@ -724,31 +729,28 @@ class WebFluxSseIntegrationTests {
assertThat(result.content().get(0)).isInstanceOf(McpSchema.TextContent.class); assertThat(result.content().get(0)).isInstanceOf(McpSchema.TextContent.class);
assertThat(((McpSchema.TextContent) result.content().get(0)).text()).isEqualTo("Logging test completed"); assertThat(((McpSchema.TextContent) result.content().get(0)).text()).isEqualTo("Logging test completed");
// Wait for notifications to be processed assertThat(latch.await(5, TimeUnit.SECONDS)).as("Should receive notifications in reasonable time").isTrue();
await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> {
// Should have received 3 notifications (1 NOTICE and 2 ERROR) // Should have received 3 notifications (1 NOTICE and 2 ERROR)
assertThat(receivedNotifications).hasSize(3); assertThat(receivedNotifications).hasSize(expectedNotificationsCount);
Map<String, McpSchema.LoggingMessageNotification> notificationMap = receivedNotifications.stream() Map<String, McpSchema.LoggingMessageNotification> notificationMap = receivedNotifications.stream()
.collect(Collectors.toMap(n -> n.data(), n -> n)); .collect(Collectors.toMap(n -> n.data(), n -> n));
// First notification should be NOTICE level // First notification should be NOTICE level
assertThat(notificationMap.get("Notice message").level()).isEqualTo(McpSchema.LoggingLevel.NOTICE); assertThat(notificationMap.get("Notice message").level()).isEqualTo(McpSchema.LoggingLevel.NOTICE);
assertThat(notificationMap.get("Notice message").logger()).isEqualTo("test-logger"); assertThat(notificationMap.get("Notice message").logger()).isEqualTo("test-logger");
assertThat(notificationMap.get("Notice message").data()).isEqualTo("Notice message"); assertThat(notificationMap.get("Notice message").data()).isEqualTo("Notice message");
// Second notification should be ERROR level // Second notification should be ERROR level
assertThat(notificationMap.get("Error message").level()).isEqualTo(McpSchema.LoggingLevel.ERROR); assertThat(notificationMap.get("Error message").level()).isEqualTo(McpSchema.LoggingLevel.ERROR);
assertThat(notificationMap.get("Error message").logger()).isEqualTo("test-logger"); assertThat(notificationMap.get("Error message").logger()).isEqualTo("test-logger");
assertThat(notificationMap.get("Error message").data()).isEqualTo("Error message"); assertThat(notificationMap.get("Error message").data()).isEqualTo("Error message");
// Third notification should be ERROR level // Third notification should be ERROR level
assertThat(notificationMap.get("Another error message").level()) assertThat(notificationMap.get("Another error message").level()).isEqualTo(McpSchema.LoggingLevel.ERROR);
.isEqualTo(McpSchema.LoggingLevel.ERROR); assertThat(notificationMap.get("Another error message").logger()).isEqualTo("test-logger");
assertThat(notificationMap.get("Another error message").logger()).isEqualTo("test-logger"); assertThat(notificationMap.get("Another error message").data()).isEqualTo("Another error message");
assertThat(notificationMap.get("Another error message").data()).isEqualTo("Another error message");
});
} }
mcpServer.close(); mcpServer.close();
} }