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

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

Browse source 
- 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>
This commit is contained in:
Luca Chang 2025-05-20 15:46:55 -07:00 • committed by Christian Tzolov
commit f3b0774a24
5 changed files with 194 additions and 37 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -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 {

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

@ -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();
});
}
}

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

@ -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);
});
}
}

67
mcp/src/test/java/io/modelcontextprotocol/client/AbstractMcpAsyncClientTests.java
View file

@ -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();
});
}
}

46
mcp/src/test/java/io/modelcontextprotocol/client/AbstractMcpSyncClientTests.java
View file

@ -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);
});
}
}