Add sampling functionality tests to MCP client test suites (#255)

- Add testSampling() methods to AbstractMcpAsyncClientTests and AbstractMcpSyncClientTests

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

Co-authored-by: Christian Tzolov <christian.tzolov@broadcom.com>

mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/client/WebClientStreamableHttpSyncClientTests.java

@@ -1,13 +1,13 @@
 package io.modelcontextprotocol.client;
 
-import com.fasterxml.jackson.databind.ObjectMapper;
 import io.modelcontextprotocol.client.transport.WebClientStreamableHttpTransport;
 import io.modelcontextprotocol.spec.McpClientTransport;
 import org.junit.jupiter.api.Timeout;
-import org.springframework.web.reactive.function.client.WebClient;
 import org.testcontainers.containers.GenericContainer;
 import org.testcontainers.containers.wait.strategy.Wait;
 
+import org.springframework.web.reactive.function.client.WebClient;
+
 @Timeout(15)
 public class WebClientStreamableHttpSyncClientTests extends AbstractMcpSyncClientTests {
 

mcp-test/src/main/java/io/modelcontextprotocol/client/AbstractMcpAsyncClientTests.java

@@ -8,6 +8,7 @@ import java.time.Duration;
 import java.util.Map;
 import java.util.Objects;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
 import java.util.function.Function;
@@ -40,6 +41,7 @@ import reactor.test.StepVerifier;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatCode;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 
 /**
  * Test suite for the {@link McpAsyncClient} that can be used with different
@@ -193,22 +195,6 @@ public abstract class AbstractMcpAsyncClientTests {
 		});
 	}
 
-	@Test
-	void testSampling() {
-		withClient(createMcpTransport(), mcpAsyncClient -> {
-			CallToolRequest callToolRequest = new CallToolRequest("sampleLLM",
-					Map.of("prompt", "Hello MCP Spring AI!"));
-			StepVerifier.create(mcpAsyncClient.initialize().then(mcpAsyncClient.callTool(callToolRequest)))
-				.consumeNextWith(callToolResult -> {
-					assertThat(callToolResult).isNotNull().satisfies(result -> {
-						assertThat(result.content()).isNotNull();
-						assertThat(result.isError()).isNull();
-					});
-				})
-				.verifyComplete();
-		});
-	}
-
 	@Test
 	void testCallToolWithInvalidTool() {
 		withClient(createMcpTransport(), mcpAsyncClient -> {
@@ -525,4 +511,52 @@ public abstract class AbstractMcpAsyncClientTests {
 		});
 	}
 
+	@Test
+	void testSampling() {
+		McpClientTransport transport = createMcpTransport();
+
+		final String message = "Hello, world!";
+		final String response = "Goodbye, world!";
+		final int maxTokens = 100;
+
+		AtomicReference<String> receivedPrompt = new AtomicReference<>();
+		AtomicReference<String> receivedMessage = new AtomicReference<>();
+		AtomicInteger receivedMaxTokens = new AtomicInteger();
+
+		withClient(transport, spec -> spec.capabilities(McpSchema.ClientCapabilities.builder().sampling().build())
+			.sampling(request -> {
+				McpSchema.TextContent messageText = assertInstanceOf(McpSchema.TextContent.class,
+						request.messages().get(0).content());
+				receivedPrompt.set(request.systemPrompt());
+				receivedMessage.set(messageText.text());
+				receivedMaxTokens.set(request.maxTokens());
+
+				return Mono
+					.just(new McpSchema.CreateMessageResult(McpSchema.Role.USER, new McpSchema.TextContent(response),
+							"modelId", McpSchema.CreateMessageResult.StopReason.END_TURN));
+			}), client -> {
+				StepVerifier.create(client.initialize()).expectNextMatches(Objects::nonNull).verifyComplete();
+
+				StepVerifier.create(client.callTool(
+						new McpSchema.CallToolRequest("sampleLLM", Map.of("prompt", message, "maxTokens", maxTokens))))
+					.consumeNextWith(result -> {
+						// Verify tool response to ensure our sampling response was passed
+						// through
+						assertThat(result.content()).hasAtLeastOneElementOfType(McpSchema.TextContent.class);
+						assertThat(result.content()).allSatisfy(content -> {
+							if (!(content instanceof McpSchema.TextContent text))
+								return;
+
+							assertThat(text.text()).endsWith(response); // Prefixed
+						});
+
+						// Verify sampling request parameters received in our callback
+						assertThat(receivedPrompt.get()).isNotEmpty();
+						assertThat(receivedMessage.get()).endsWith(message); // Prefixed
+						assertThat(receivedMaxTokens.get()).isEqualTo(maxTokens);
+					})
+					.verifyComplete();
+			});
+	}
+
 }

mcp-test/src/main/java/io/modelcontextprotocol/client/AbstractMcpSyncClientTests.java

@@ -8,6 +8,7 @@ import java.time.Duration;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
 import java.util.function.Function;
@@ -31,13 +32,12 @@ import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import reactor.core.publisher.Mono;
-import reactor.core.scheduler.Scheduler;
-import reactor.core.scheduler.Schedulers;
 import reactor.test.StepVerifier;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatCode;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 
 /**
  * Unit tests for MCP Client Session functionality.
@@ -438,4 +438,48 @@ public abstract class AbstractMcpSyncClientTests {
 			.hasMessageContaining("Logging level must not be null"));
 	}
 
+	@Test
+	void testSampling() {
+		McpClientTransport transport = createMcpTransport();
+
+		final String message = "Hello, world!";
+		final String response = "Goodbye, world!";
+		final int maxTokens = 100;
+
+		AtomicReference<String> receivedPrompt = new AtomicReference<>();
+		AtomicReference<String> receivedMessage = new AtomicReference<>();
+		AtomicInteger receivedMaxTokens = new AtomicInteger();
+
+		withClient(transport, spec -> spec.capabilities(McpSchema.ClientCapabilities.builder().sampling().build())
+			.sampling(request -> {
+				McpSchema.TextContent messageText = assertInstanceOf(McpSchema.TextContent.class,
+						request.messages().get(0).content());
+				receivedPrompt.set(request.systemPrompt());
+				receivedMessage.set(messageText.text());
+				receivedMaxTokens.set(request.maxTokens());
+
+				return new McpSchema.CreateMessageResult(McpSchema.Role.USER, new McpSchema.TextContent(response),
+						"modelId", McpSchema.CreateMessageResult.StopReason.END_TURN);
+			}), client -> {
+				client.initialize();
+
+				McpSchema.CallToolResult result = client.callTool(
+						new McpSchema.CallToolRequest("sampleLLM", Map.of("prompt", message, "maxTokens", maxTokens)));
+
+				// Verify tool response to ensure our sampling response was passed through
+				assertThat(result.content()).hasAtLeastOneElementOfType(McpSchema.TextContent.class);
+				assertThat(result.content()).allSatisfy(content -> {
+					if (!(content instanceof McpSchema.TextContent text))
+						return;
+
+					assertThat(text.text()).endsWith(response); // Prefixed
+				});
+
+				// Verify sampling request parameters received in our callback
+				assertThat(receivedPrompt.get()).isNotEmpty();
+				assertThat(receivedMessage.get()).endsWith(message); // Prefixed
+				assertThat(receivedMaxTokens.get()).isEqualTo(maxTokens);
+			});
+	}
+
 }

mcp/src/test/java/io/modelcontextprotocol/client/AbstractMcpAsyncClientTests.java

@@ -8,6 +8,7 @@ import java.time.Duration;
 import java.util.Map;
 import java.util.Objects;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
 import java.util.function.Function;
@@ -40,6 +41,7 @@ import reactor.test.StepVerifier;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatCode;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 
 /**
  * Test suite for the {@link McpAsyncClient} that can be used with different
@@ -194,22 +196,6 @@ public abstract class AbstractMcpAsyncClientTests {
 		});
 	}
 
-	@Test
-	void testSampling() {
-		withClient(createMcpTransport(), mcpAsyncClient -> {
-			CallToolRequest callToolRequest = new CallToolRequest("sampleLLM",
-					Map.of("prompt", "Hello MCP Spring AI!"));
-			StepVerifier.create(mcpAsyncClient.initialize().then(mcpAsyncClient.callTool(callToolRequest)))
-				.consumeNextWith(callToolResult -> {
-					assertThat(callToolResult).isNotNull().satisfies(result -> {
-						assertThat(result.content()).isNotNull();
-						assertThat(result.isError()).isNull();
-					});
-				})
-				.verifyComplete();
-		});
-	}
-
 	@Test
 	void testCallToolWithInvalidTool() {
 		withClient(createMcpTransport(), mcpAsyncClient -> {
@@ -486,7 +472,6 @@ public abstract class AbstractMcpAsyncClientTests {
 	// ---------------------------------------
 	// Logging Tests
 	// ---------------------------------------
-
 	@Test
 	void testLoggingLevelsWithoutInitialization() {
 		verifyNotificationSucceedsWithImplicitInitialization(
@@ -526,4 +511,52 @@ public abstract class AbstractMcpAsyncClientTests {
 		});
 	}
 
+	@Test
+	void testSampling() {
+		McpClientTransport transport = createMcpTransport();
+
+		final String message = "Hello, world!";
+		final String response = "Goodbye, world!";
+		final int maxTokens = 100;
+
+		AtomicReference<String> receivedPrompt = new AtomicReference<>();
+		AtomicReference<String> receivedMessage = new AtomicReference<>();
+		AtomicInteger receivedMaxTokens = new AtomicInteger();
+
+		withClient(transport, spec -> spec.capabilities(McpSchema.ClientCapabilities.builder().sampling().build())
+			.sampling(request -> {
+				McpSchema.TextContent messageText = assertInstanceOf(McpSchema.TextContent.class,
+						request.messages().get(0).content());
+				receivedPrompt.set(request.systemPrompt());
+				receivedMessage.set(messageText.text());
+				receivedMaxTokens.set(request.maxTokens());
+
+				return Mono
+					.just(new McpSchema.CreateMessageResult(McpSchema.Role.USER, new McpSchema.TextContent(response),
+							"modelId", McpSchema.CreateMessageResult.StopReason.END_TURN));
+			}), client -> {
+				StepVerifier.create(client.initialize()).expectNextMatches(Objects::nonNull).verifyComplete();
+
+				StepVerifier.create(client.callTool(
+						new McpSchema.CallToolRequest("sampleLLM", Map.of("prompt", message, "maxTokens", maxTokens))))
+					.consumeNextWith(result -> {
+						// Verify tool response to ensure our sampling response was passed
+						// through
+						assertThat(result.content()).hasAtLeastOneElementOfType(McpSchema.TextContent.class);
+						assertThat(result.content()).allSatisfy(content -> {
+							if (!(content instanceof McpSchema.TextContent text))
+								return;
+
+							assertThat(text.text()).endsWith(response); // Prefixed
+						});
+
+						// Verify sampling request parameters received in our callback
+						assertThat(receivedPrompt.get()).isNotEmpty();
+						assertThat(receivedMessage.get()).endsWith(message); // Prefixed
+						assertThat(receivedMaxTokens.get()).isEqualTo(maxTokens);
+					})
+					.verifyComplete();
+			});
+	}
+
 }

mcp/src/test/java/io/modelcontextprotocol/client/AbstractMcpSyncClientTests.java

@@ -8,6 +8,7 @@ import java.time.Duration;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
 import java.util.function.Function;
@@ -38,6 +39,7 @@ import reactor.test.StepVerifier;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatCode;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 
 /**
  * Unit tests for MCP Client Session functionality.
@@ -439,4 +441,48 @@ public abstract class AbstractMcpSyncClientTests {
 			.hasMessageContaining("Logging level must not be null"));
 	}
 
+	@Test
+	void testSampling() {
+		McpClientTransport transport = createMcpTransport();
+
+		final String message = "Hello, world!";
+		final String response = "Goodbye, world!";
+		final int maxTokens = 100;
+
+		AtomicReference<String> receivedPrompt = new AtomicReference<>();
+		AtomicReference<String> receivedMessage = new AtomicReference<>();
+		AtomicInteger receivedMaxTokens = new AtomicInteger();
+
+		withClient(transport, spec -> spec.capabilities(McpSchema.ClientCapabilities.builder().sampling().build())
+			.sampling(request -> {
+				McpSchema.TextContent messageText = assertInstanceOf(McpSchema.TextContent.class,
+						request.messages().get(0).content());
+				receivedPrompt.set(request.systemPrompt());
+				receivedMessage.set(messageText.text());
+				receivedMaxTokens.set(request.maxTokens());
+
+				return new McpSchema.CreateMessageResult(McpSchema.Role.USER, new McpSchema.TextContent(response),
+						"modelId", McpSchema.CreateMessageResult.StopReason.END_TURN);
+			}), client -> {
+				client.initialize();
+
+				McpSchema.CallToolResult result = client.callTool(
+						new McpSchema.CallToolRequest("sampleLLM", Map.of("prompt", message, "maxTokens", maxTokens)));
+
+				// Verify tool response to ensure our sampling response was passed through
+				assertThat(result.content()).hasAtLeastOneElementOfType(McpSchema.TextContent.class);
+				assertThat(result.content()).allSatisfy(content -> {
+					if (!(content instanceof McpSchema.TextContent text))
+						return;
+
+					assertThat(text.text()).endsWith(response); // Prefixed
+				});
+
+				// Verify sampling request parameters received in our callback
+				assertThat(receivedPrompt.get()).isNotEmpty();
+				assertThat(receivedMessage.get()).endsWith(message); // Prefixed
+				assertThat(receivedMaxTokens.get()).isEqualTo(maxTokens);
+			});
+	}
+
 }