Clean the code and add more tests

This commit is contained in:
Christian Tzolov 2024-12-09 23:05:23 +01:00
commit 68c21d2264
8 changed files with 312 additions and 732 deletions

View file

@ -25,7 +25,7 @@ import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.Timeout;
import spring.ai.mcp.client.McpAsyncClient;
import spring.ai.mcp.client.McpClient;
import spring.ai.mcp.client.stdio.StdioServerParameters;
import spring.ai.mcp.client.stdio.ServerParameters;
import spring.ai.mcp.client.stdio.StdioServerTransport;
import spring.ai.mcp.spec.McpSchema.CallToolRequest;
import spring.ai.mcp.spec.McpSchema.Resource;
@ -45,7 +45,7 @@ class McpAsyncClientTests {
private McpAsyncClient mcpAsyncClient;
private StdioServerParameters stdioParams;
private ServerParameters stdioParams;
private static final Duration TIMEOUT = Duration.ofSeconds(10);
@ -53,7 +53,7 @@ class McpAsyncClientTests {
@BeforeEach
void setUp() {
stdioParams = StdioServerParameters.builder("npx")
stdioParams = ServerParameters.builder("npx")
.args("-y", "@modelcontextprotocol/server-everything", "dir")
.build();

View file

@ -16,6 +16,7 @@
package spring.ai.mcp;
import java.time.Duration;
import java.util.List;
import java.util.Map;
import org.junit.jupiter.api.AfterEach;
@ -28,14 +29,16 @@ import static org.assertj.core.api.Assertions.*;
import com.fasterxml.jackson.databind.ObjectMapper;
import spring.ai.mcp.client.McpClient;
import spring.ai.mcp.client.McpSyncClient;
import spring.ai.mcp.client.stdio.StdioServerParameters;
import spring.ai.mcp.client.stdio.ServerParameters;
import spring.ai.mcp.client.stdio.StdioServerTransport;
import spring.ai.mcp.spec.McpSchema.CallToolRequest;
import spring.ai.mcp.spec.McpSchema.CallToolResult;
import spring.ai.mcp.spec.McpSchema.Content;
import spring.ai.mcp.spec.McpSchema.ListResourcesResult;
import spring.ai.mcp.spec.McpSchema.ListToolsResult;
import spring.ai.mcp.spec.McpSchema.Tool;
import spring.ai.mcp.spec.McpSchema.Resource;
import spring.ai.mcp.spec.McpSchema.TextContent;
/**
* Unit tests for MCP Client Session functionality.
@ -47,7 +50,7 @@ class McpSyncClientTests {
private McpSyncClient mcpSyncClient;
private StdioServerParameters stdioParams;
private ServerParameters stdioParams;
private static final Duration TIMEOUT = Duration.ofSeconds(10);
@ -55,7 +58,7 @@ class McpSyncClientTests {
@BeforeEach
void setUp() {
stdioParams = StdioServerParameters.builder("npx")
stdioParams = ServerParameters.builder("npx")
.args("-y", "@modelcontextprotocol/server-everything", "dir")
.build();
@ -98,6 +101,25 @@ class McpSyncClientTests {
});
}
@Test
@Timeout(15) // Giving extra time beyond the client timeout
void testCallTools() {
CallToolResult toolResult = mcpSyncClient.callTool(new CallToolRequest("add", Map.of("a", 3, "b", 4)));
assertThat(toolResult).isNotNull().satisfies(result -> {
assertThat(result.content()).hasSize(1);
TextContent content = (TextContent) result.content().get(0);
assertThat(content).isNotNull();
assertThat(content.text()).isNotNull();
assertThat(content.text()).contains("7");
System.out.println(content);
});
}
@Test
@Timeout(15)
void testPing() {

View file

@ -20,7 +20,7 @@ import java.time.Duration;
import com.fasterxml.jackson.databind.ObjectMapper;
import spring.ai.mcp.client.McpClient;
import spring.ai.mcp.client.McpSyncClient;
import spring.ai.mcp.client.stdio.StdioServerParameters;
import spring.ai.mcp.client.stdio.ServerParameters;
import spring.ai.mcp.client.stdio.StdioServerTransport;
import spring.ai.mcp.spec.McpSchema.ListResourcesResult;
import spring.ai.mcp.spec.McpSchema.ListToolsResult;
@ -34,7 +34,7 @@ public class ClientSessionTests2 {
public static void main(String[] args) {
var stdioParams = StdioServerParameters.builder("uv")
var stdioParams = ServerParameters.builder("uv")
.args("--directory", "dir", "run", "mcp-server-sqlite",
"--db-path", "~/test.db")
.build();

View file

@ -20,7 +20,7 @@ import java.time.Duration;
import com.fasterxml.jackson.databind.ObjectMapper;
import spring.ai.mcp.client.McpClient;
import spring.ai.mcp.client.McpSyncClient;
import spring.ai.mcp.client.stdio.StdioServerParameters;
import spring.ai.mcp.client.stdio.ServerParameters;
import spring.ai.mcp.client.stdio.StdioServerTransport;
import spring.ai.mcp.spec.McpSchema.ListResourcesResult;
import spring.ai.mcp.spec.McpSchema.ListToolsResult;
@ -34,7 +34,7 @@ public class ClientSessionTests3 {
public static void main(String[] args) {
var stdioParams = StdioServerParameters.builder("uv")
var stdioParams = ServerParameters.builder("uv")
.args("--directory", "dir", "run", "mcp-server-sqlite",
"--db-path", "~/test.db")
.build();

View file

@ -0,0 +1,161 @@
/*
* Copyright 2024 - 2024 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package spring.ai.mcp.spec;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import reactor.core.publisher.Mono;
import reactor.test.StepVerifier;
import spring.ai.mcp.spec.McpSchema.JSONRPCMessage;
import spring.ai.mcp.spec.McpSchema.JSONRPCRequest;
import java.util.concurrent.atomic.AtomicReference;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
class DefaultMcpTransportTests {
private DefaultMcpTransport transport;
private ObjectMapper objectMapper;
@BeforeEach
void setUp() {
objectMapper = new ObjectMapper();
transport = new DefaultMcpTransport(objectMapper);
}
@Test
void constructorShouldValidateObjectMapper() {
assertThatThrownBy(() -> new DefaultMcpTransport(null))
.isInstanceOf(IllegalArgumentException.class)
.hasMessage("ObjectMapper must not be null");
}
@Test
void defaultConstructorShouldCreateValidInstance() {
DefaultMcpTransport defaultTransport = new DefaultMcpTransport();
assertThat(defaultTransport.getObjectMapper()).isNotNull();
}
@Test
void customObjectMapperConstructorShouldCreateValidInstance() {
assertThat(transport.getObjectMapper()).isSameAs(objectMapper);
}
@Test
void sendMessageShouldEmitToOutboundSink() {
JSONRPCRequest message = new JSONRPCRequest(
"2.0",
"test",
1,
null);
AtomicReference<JSONRPCMessage> receivedMessage = new AtomicReference<>();
transport.getOutboundSink().asFlux().subscribe(receivedMessage::set);
Mono<Void> result = transport.sendMessage(message);
StepVerifier.create(result)
.verifyComplete();
assertThat(receivedMessage.get())
.isNotNull()
.satisfies(msg -> {
assertThat(((JSONRPCRequest) msg).jsonrpc()).isEqualTo("2.0");
assertThat(((JSONRPCRequest) msg).method()).isEqualTo("test");
});
}
@Test
void customInboundMessageHandlerShouldReceiveMessages() {
AtomicReference<JSONRPCMessage> receivedMessage = new AtomicReference<>();
transport.setInboudMessageHandler(receivedMessage::set);
transport.start();
JSONRPCRequest message = new JSONRPCRequest(
"2.0",
"test",
1,
null);
transport.getInboundSink().tryEmitNext(message);
assertThat(receivedMessage.get())
.isNotNull()
.satisfies(msg -> {
assertThat(((JSONRPCRequest) msg).jsonrpc()).isEqualTo("2.0");
assertThat(((JSONRPCRequest) msg).method()).isEqualTo("test");
});
}
@Test
void customErrorHandlerShouldReceiveErrors() {
AtomicReference<String> receivedError = new AtomicReference<>();
transport.setErrorHandler(receivedError::set);
transport.start();
String errorMessage = "Test error";
transport.getErrorSink().tryEmitNext(errorMessage);
assertThat(receivedError.get())
.isNotNull()
.isEqualTo(errorMessage);
}
@Test
void startShouldInitializeMessageAndErrorHandling() {
AtomicReference<JSONRPCMessage> receivedMessage = new AtomicReference<>();
AtomicReference<String> receivedError = new AtomicReference<>();
transport.setInboudMessageHandler(receivedMessage::set);
transport.setErrorHandler(receivedError::set);
transport.start();
JSONRPCRequest message = new JSONRPCRequest(
"2.0",
"test",
1,
null);
String errorMessage = "Test error";
transport.getInboundSink().tryEmitNext(message);
transport.getErrorSink().tryEmitNext(errorMessage);
assertThat(receivedMessage.get()).isNotNull();
assertThat(receivedError.get()).isEqualTo(errorMessage);
}
@Test
void sendMessageShouldReturnErrorWhenSinkFails() {
JSONRPCRequest message = new JSONRPCRequest(
"2.0",
"test",
1,
null);
// Close the sink to simulate failure
transport.getOutboundSink().tryEmitComplete();
Mono<Void> result = transport.sendMessage(message);
StepVerifier.create(result)
.verifyErrorSatisfies(error -> assertThat(error)
.isInstanceOf(RuntimeException.class)
.hasMessage("Failed to enqueue message"));
}
}

View file

@ -1,158 +0,0 @@
/*
* Copyright 2024 - 2024 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package spring.ai.mcp.spec;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Disabled;
import org.junit.jupiter.api.Test;
import org.mockito.ArgumentCaptor;
import spring.ai.mcp.spec.McpSchema.JSONRPCMessage;
import spring.ai.mcp.spec.McpSchema.JSONRPCNotification;
import spring.ai.mcp.spec.McpSchema.JSONRPCRequest;
import spring.ai.mcp.spec.McpSchema.JSONRPCResponse;
import java.time.Duration;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.function.Consumer;
import static org.junit.jupiter.api.Assertions.*;
import static org.mockito.Mockito.*;
class McpSessionJSONRPCTest {
private McpTransport mockTransport;
private ObjectMapper objectMapper;
private McpSession mcpSession;
private ArgumentCaptor<Consumer<JSONRPCMessage>> messageHandlerCaptor;
@BeforeEach
void setUp() {
mockTransport = mock(McpTransport.class);
objectMapper = new ObjectMapper();
mcpSession = new McpSession(mockTransport, Duration.ofSeconds(5), objectMapper);
// Capture the message handler set during the session initialization
messageHandlerCaptor = ArgumentCaptor.forClass(Consumer.class);
verify(mockTransport).setMessageHandler(messageHandlerCaptor.capture());
}
@Test
void sendRequestShouldSendValidRequestMessage() {
String method = "testMethod";
Map<String, Object> params = Map.of("key", "value");
// Call sendRequest
CompletableFuture<String> future = mcpSession.sendRequest(method, params, new TypeReference<>() {});
// Capture the JSONRPCRequest sent to the transport
ArgumentCaptor<JSONRPCMessage> messageCaptor = ArgumentCaptor.forClass(JSONRPCMessage.class);
verify(mockTransport).sendMessage(messageCaptor.capture());
// Assert the message details
JSONRPCRequest request = (JSONRPCRequest) messageCaptor.getValue();
assertEquals(method, request.method());
assertEquals("value", ((Map<?, ?>) request.params()).get("key"));
assertNotNull(request.id());
}
@Test
void sendNotificationShouldSendValidNotificationMessage() {
String method = "notifyMethod";
Map<String, Object> params = Map.of("param1", "value1");
// Call sendNotification
mcpSession.sendNotification(method, params);
// Capture the JSONRPCNotification sent to the transport
ArgumentCaptor<JSONRPCMessage> messageCaptor = ArgumentCaptor.forClass(JSONRPCMessage.class);
verify(mockTransport).sendMessage(messageCaptor.capture());
// Assert the notification details
JSONRPCNotification notification = (JSONRPCNotification) messageCaptor.getValue();
assertEquals(method, notification.method());
assertEquals("value1", ((Map<?, ?>) notification.params()).get("param1"));
}
@Test
void messageHandlerShouldProcessResponseCorrectly() {
// Retrieve the captured message handler
Consumer<JSONRPCMessage> messageHandler = messageHandlerCaptor.getValue();
// Simulate a JSONRPCResponse
JSONRPCResponse response = new JSONRPCResponse("2.0", "1", "resultValue", null);
CompletableFuture<JSONRPCResponse> responseFuture = new CompletableFuture<>();
mcpSession.waitingForResponseStreams.put("1", responseFuture);
// Trigger the handler with the response
messageHandler.accept(response);
// Verify the response is completed
assertTrue(responseFuture.isDone());
assertEquals("resultValue", responseFuture.join().result());
}
@Test
@Disabled
void messageHandlerShouldInvokeHandleRequestForRequestMessage() {
// Retrieve the captured message handler
Consumer<JSONRPCMessage> messageHandler = messageHandlerCaptor.getValue();
// Mock request data
JSONRPCRequest request = new JSONRPCRequest("2.0", "testMethod", "1", Map.of("key", "value"));
// Override the handleRequest to check invocation
McpSession spySession = spy(mcpSession);
doNothing().when(spySession).handleRequest(any());
// Trigger the handler with a request
messageHandler.accept(request);
// Verify handleRequest was called
verify(spySession).handleRequest(any());
}
@Test
@Disabled
void messageHandlerShouldInvokeHandleNotificationForNotificationMessage() {
// Retrieve the captured message handler
Consumer<JSONRPCMessage> messageHandler = messageHandlerCaptor.getValue();
// Mock notification data
JSONRPCNotification notification = new JSONRPCNotification("2.0", "notifyMethod", Map.of("key", "value"));
// Override the handleNotification to check invocation
McpSession spySession = spy(mcpSession);
doNothing().when(spySession).handleNotification(any());
// Trigger the handler with a notification
messageHandler.accept(notification);
// Verify handleNotification was called
verify(spySession).handleNotification(notification);
}
@Test
void closeShouldInvokeTransportClose() {
// Call close on the session
mcpSession.close();
// Verify transport close was called
verify(mockTransport).close();
}
}

View file

@ -1,93 +0,0 @@
/*
* Copyright 2024 - 2024 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package spring.ai.mcp.spec;
import java.time.Duration;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Disabled;
import org.junit.jupiter.api.Test;
import spring.ai.mcp.spec.McpSchema.JSONRPCResponse;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
class McpSessionTest {
private McpTransport mockTransport;
private ObjectMapper mockObjectMapper;
private McpSession mcpSession;
@BeforeEach
void setUp() {
mockTransport = mock(McpTransport.class);
mockObjectMapper = mock(ObjectMapper.class);
mcpSession = new McpSession(mockTransport, Duration.ofSeconds(5), mockObjectMapper);
}
@Test
void constructorShouldThrowIfArgumentsAreNull() {
assertThrows(IllegalArgumentException.class,
() -> new McpSession(null, Duration.ofSeconds(5), mockObjectMapper));
assertThrows(IllegalArgumentException.class, () -> new McpSession(mockTransport, null, mockObjectMapper));
assertThrows(IllegalArgumentException.class, () -> new McpSession(mockTransport, Duration.ofSeconds(5), null));
}
@Test
@Disabled
void sendRequestShouldSendJsonRpcRequest() throws Exception {
String method = "testMethod";
Object params = Map.of("key", "value");
TypeReference<String> typeRef = new TypeReference<>() {
};
CompletableFuture<McpSchema.JSONRPCResponse> mockResponseFuture = new CompletableFuture<>();
mockResponseFuture.complete(new JSONRPCResponse("2.0", "1", "result", null));
// when(mockTransport.sendMessage(any())).then(invocation -> null);
when(mockObjectMapper.convertValue("result", typeRef)).thenReturn("result");
CompletableFuture<String> result = mcpSession.sendRequest(method, params, typeRef);
assertEquals("result", result.get());
}
@Test
void sendRequestShouldHandleTimeout() {
String method = "timeoutMethod";
Object params = Map.of("key", "value");
TypeReference<String> typeRef = new TypeReference<>() {
};
CompletableFuture<String> result = mcpSession.sendRequest(method, params, typeRef);
assertThrows(Exception.class, result::get);
}
@Test
void closeShouldInvokeTransportClose() {
mcpSession.close();
verify(mockTransport, times(1)).close();
}
}

View file

@ -1,128 +0,0 @@
/*
* Copyright 2024 - 2024 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package spring.ai.mcp.spec;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.mockito.ArgumentCaptor;
import spring.ai.mcp.spec.McpSchema.JSONRPCMessage;
import spring.ai.mcp.spec.McpSchema.JSONRPCNotification;
import spring.ai.mcp.spec.McpSchema.JSONRPCRequest;
import spring.ai.mcp.spec.McpSchema.JSONRPCResponse;
import java.time.Duration;
import java.util.Map;
import java.util.function.Consumer;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import static org.junit.jupiter.api.Assertions.*;
import static org.mockito.Mockito.*;
class McpTransportTest {
private McpTransport mockTransport;
private McpSession mcpSession;
@BeforeEach
void setUp() {
mockTransport = mock(McpTransport.class);
mcpSession = new McpSession(mockTransport, Duration.ofSeconds(5), new ObjectMapper());
}
@Test
void constructorShouldInitializeTransport() {
// Verify transport setup during construction
verify(mockTransport).setMessageHandler(any(Consumer.class));
verify(mockTransport).setErrorHandler(any(Consumer.class));
verify(mockTransport).start();
}
@Test
void sendRequestShouldCallTransportSendMessage() {
String method = "testMethod";
Object params = Map.of("key", "value");
// Call sendRequest
mcpSession.sendRequest(method, params, new TypeReference<String>() {});
// Capture the message sent to transport
ArgumentCaptor<JSONRPCMessage> messageCaptor = ArgumentCaptor.forClass(JSONRPCMessage.class);
verify(mockTransport).sendMessage(messageCaptor.capture());
// Verify the message details
JSONRPCRequest request = (JSONRPCRequest) messageCaptor.getValue();
assertEquals("testMethod", request.method());
assertEquals("value", ((Map<?, ?>) request.params()).get("key"));
assertNotNull(request.id());
}
@Test
void sendNotificationShouldCallTransportSendMessage() {
String method = "notifyMethod";
Map<String, Object> params = Map.of("param1", "value1");
// Call sendNotification
mcpSession.sendNotification(method, params);
// Capture the notification sent to transport
ArgumentCaptor<JSONRPCMessage> messageCaptor = ArgumentCaptor.forClass(JSONRPCMessage.class);
verify(mockTransport).sendMessage(messageCaptor.capture());
// Verify the notification details
JSONRPCNotification notification = (JSONRPCNotification) messageCaptor.getValue();
assertEquals("notifyMethod", notification.method());
assertEquals("value1", ((Map<?, ?>) notification.params()).get("param1"));
}
@Test
void messageHandlerShouldProcessResponse() {
// Capture the message handler during initialization
ArgumentCaptor<Consumer<JSONRPCMessage>> handlerCaptor = ArgumentCaptor.forClass(Consumer.class);
verify(mockTransport).setMessageHandler(handlerCaptor.capture());
Consumer<JSONRPCMessage> handler = handlerCaptor.getValue();
// Simulate a JSONRPCResponse
JSONRPCResponse response = new JSONRPCResponse("2.0", "1", "result", null);
handler.accept(response);
// Verify the response was processed
// (Use CompletableFuture to match the response if needed in real tests)
}
@Test
void errorHandlerShouldLogErrors() {
// Capture the error handler during initialization
ArgumentCaptor<Consumer<String>> handlerCaptor = ArgumentCaptor.forClass(Consumer.class);
verify(mockTransport).setErrorHandler(handlerCaptor.capture());
Consumer<String> errorHandler = handlerCaptor.getValue();
// Simulate an error
String errorMessage = "Test error";
errorHandler.accept(errorMessage);
// Verify the error was logged or processed
// (In real tests, capture logging output or other side effects)
}
@Test
void closeShouldStopTransport() {
mcpSession.close();
verify(mockTransport).close();
}
}