Skip to content

Commit 2f87aed

Browse files
tzolovtzolov
committed
Address review comments and improve the test run isolation
Signed-off-by: Christian Tzolov <christian.tzolov@broadcom.com>
1 parent 09d8df4 commit 2f87aed

File tree

1 file changed

+66
-88
lines changed

1 file changed

+66
-88
lines changed

‎mcp-core/src/test/java/io/modelcontextprotocol/spec/McpClientSessionTests.java‎

Lines changed: 66 additions & 88 deletions
Original file line numberDiff line numberDiff line change
@@ -6,11 +6,10 @@
66

77
import java.time.Duration;
88
import java.util.Map;
9+
import java.util.function.Function;
910

1011
import io.modelcontextprotocol.MockMcpClientTransport;
1112
import io.modelcontextprotocol.json.TypeRef;
12-
import org.junit.jupiter.api.AfterEach;
13-
import org.junit.jupiter.api.BeforeEach;
1413
import org.junit.jupiter.api.Test;
1514
import org.slf4j.Logger;
1615
import org.slf4j.LoggerFactory;
@@ -38,24 +37,6 @@ class McpClientSessionTests {
3837

3938
private static final String ECHO_METHOD = "echo";
4039

41-
private McpClientSession session;
42-
43-
private MockMcpClientTransport transport;
44-
45-
@BeforeEach
46-
void setUp() {
47-
transport = new MockMcpClientTransport();
48-
session = new McpClientSession(TIMEOUT, transport, Map.of(),
49-
Map.of(TEST_NOTIFICATION, params -> Mono.fromRunnable(() -> logger.info("Status update: {}", params))));
50-
}
51-
52-
@AfterEach
53-
void tearDown() {
54-
if (session != null) {
55-
session.close();
56-
}
57-
}
58-
5940
TypeRef<String> responseType = new TypeRef<>() {
6041
};
6142

@@ -64,6 +45,11 @@ void testSendRequest() {
6445
String testParam = "test parameter";
6546
String responseData = "test response";
6647

48+
var transport = new MockMcpClientTransport();
49+
var session = new McpClientSession(TIMEOUT, transport, Map.of(),
50+
Map.of(TEST_NOTIFICATION, params -> Mono.fromRunnable(() -> logger.info("Status update: {}", params))),
51+
Function.identity());
52+
6753
// Create a Mono that will emit the response after the request is sent
6854
Mono<String> responseMono = session.sendRequest(TEST_METHOD, testParam, responseType);
6955
// Verify response handling
@@ -80,10 +66,17 @@ void testSendRequest() {
8066
assertThat(request.params()).isEqualTo(testParam);
8167
assertThat(response).isEqualTo(responseData);
8268
}).verifyComplete();
69+
70+
session.close();
8371
}
8472

8573
@Test
8674
void testSendRequestWithError() {
75+
var transport = new MockMcpClientTransport();
76+
var session = new McpClientSession(TIMEOUT, transport, Map.of(),
77+
Map.of(TEST_NOTIFICATION, params -> Mono.fromRunnable(() -> logger.info("Status update: {}", params))),
78+
Function.identity());
79+
8780
Mono<String> responseMono = session.sendRequest(TEST_METHOD, "test", responseType);
8881

8982
// Verify error handling
@@ -95,20 +88,34 @@ void testSendRequestWithError() {
9588
transport.simulateIncomingMessage(
9689
new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), null, error));
9790
}).expectError(McpError.class).verify();
91+
92+
session.close();
9893
}
9994

10095
@Test
10196
void testRequestTimeout() {
97+
var transport = new MockMcpClientTransport();
98+
var session = new McpClientSession(TIMEOUT, transport, Map.of(),
99+
Map.of(TEST_NOTIFICATION, params -> Mono.fromRunnable(() -> logger.info("Status update: {}", params))),
100+
Function.identity());
101+
102102
Mono<String> responseMono = session.sendRequest(TEST_METHOD, "test", responseType);
103103

104104
// Verify timeout
105105
StepVerifier.create(responseMono)
106106
.expectError(java.util.concurrent.TimeoutException.class)
107107
.verify(TIMEOUT.plusSeconds(1));
108+
109+
session.close();
108110
}
109111

110112
@Test
111113
void testSendNotification() {
114+
var transport = new MockMcpClientTransport();
115+
var session = new McpClientSession(TIMEOUT, transport, Map.of(),
116+
Map.of(TEST_NOTIFICATION, params -> Mono.fromRunnable(() -> logger.info("Status update: {}", params))),
117+
Function.identity());
118+
112119
Map<String, Object> params = Map.of("key", "value");
113120
Mono<Void> notificationMono = session.sendNotification(TEST_NOTIFICATION, params);
114121

@@ -120,15 +127,17 @@ void testSendNotification() {
120127
assertThat(notification.method()).isEqualTo(TEST_NOTIFICATION);
121128
assertThat(notification.params()).isEqualTo(params);
122129
}).verifyComplete();
130+
131+
session.close();
123132
}
124133

125134
@Test
126135
void testRequestHandling() {
127136
String echoMessage = "Hello MCP!";
128137
Map<String, McpClientSession.RequestHandler<?>> requestHandlers = Map.of(ECHO_METHOD,
129138
params -> Mono.just(params));
130-
transport = new MockMcpClientTransport();
131-
session = new McpClientSession(TIMEOUT, transport, requestHandlers, Map.of());
139+
var transport = new MockMcpClientTransport();
140+
var session = new McpClientSession(TIMEOUT, transport, requestHandlers, Map.of(), Function.identity());
132141

133142
// Simulate incoming request
134143
McpSchema.JSONRPCRequest request = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, ECHO_METHOD,
@@ -141,15 +150,18 @@ void testRequestHandling() {
141150
McpSchema.JSONRPCResponse response = (McpSchema.JSONRPCResponse) sentMessage;
142151
assertThat(response.result()).isEqualTo(echoMessage);
143152
assertThat(response.error()).isNull();
153+
154+
session.close();
144155
}
145156

146157
@Test
147158
void testNotificationHandling() {
148159
Sinks.One<Object> receivedParams = Sinks.one();
149160

150-
transport = new MockMcpClientTransport();
151-
session = new McpClientSession(TIMEOUT, transport, Map.of(),
152-
Map.of(TEST_NOTIFICATION, params -> Mono.fromRunnable(() -> receivedParams.tryEmitValue(params))));
161+
var transport = new MockMcpClientTransport();
162+
var session = new McpClientSession(TIMEOUT, transport, Map.of(),
163+
Map.of(TEST_NOTIFICATION, params -> Mono.fromRunnable(() -> receivedParams.tryEmitValue(params))),
164+
Function.identity());
153165

154166
// Simulate incoming notification from the server
155167
Map<String, Object> notificationParams = Map.of("status", "ready");
@@ -161,10 +173,18 @@ void testNotificationHandling() {
161173

162174
// Verify handler was called
163175
assertThat(receivedParams.asMono().block(Duration.ofSeconds(1))).isEqualTo(notificationParams);
176+
177+
session.close();
164178
}
165179

166180
@Test
167181
void testUnknownMethodHandling() {
182+
183+
var transport = new MockMcpClientTransport();
184+
var session = new McpClientSession(TIMEOUT, transport, Map.of(),
185+
Map.of(TEST_NOTIFICATION, params -> Mono.fromRunnable(() -> logger.info("Status update: {}", params))),
186+
Function.identity());
187+
168188
// Simulate incoming request for unknown method
169189
McpSchema.JSONRPCRequest request = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, "unknown.method",
170190
"test-id", null);
@@ -176,6 +196,8 @@ void testUnknownMethodHandling() {
176196
McpSchema.JSONRPCResponse response = (McpSchema.JSONRPCResponse) sentMessage;
177197
assertThat(response.error()).isNotNull();
178198
assertThat(response.error().code()).isEqualTo(McpSchema.ErrorCodes.METHOD_NOT_FOUND);
199+
200+
session.close();
179201
}
180202

181203
@Test
@@ -187,8 +209,9 @@ void testRequestHandlerThrowsMcpErrorWithJsonRpcError() {
187209
McpClientSession.RequestHandler<?> failingHandler = params -> Mono
188210
.error(McpError.builder(123).message("Custom error message").data(errorData).build());
189211

190-
transport = new MockMcpClientTransport();
191-
session = new McpClientSession(TIMEOUT, transport, Map.of(testMethod, failingHandler), Map.of());
212+
var transport = new MockMcpClientTransport();
213+
var session = new McpClientSession(TIMEOUT, transport, Map.of(testMethod, failingHandler), Map.of(),
214+
Function.identity());
192215

193216
// Simulate incoming request that will trigger the error
194217
McpSchema.JSONRPCRequest request = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, testMethod,
@@ -203,6 +226,8 @@ void testRequestHandlerThrowsMcpErrorWithJsonRpcError() {
203226
assertThat(response.error().code()).isEqualTo(123);
204227
assertThat(response.error().message()).isEqualTo("Custom error message");
205228
assertThat(response.error().data()).isEqualTo(errorData);
229+
230+
session.close();
206231
}
207232

208233
@Test
@@ -212,8 +237,9 @@ void testRequestHandlerThrowsGenericException() {
212237
RuntimeException exception = new RuntimeException("Something went wrong");
213238
McpClientSession.RequestHandler<?> failingHandler = params -> Mono.error(exception);
214239

215-
transport = new MockMcpClientTransport();
216-
session = new McpClientSession(TIMEOUT, transport, Map.of(testMethod, failingHandler), Map.of());
240+
var transport = new MockMcpClientTransport();
241+
var session = new McpClientSession(TIMEOUT, transport, Map.of(testMethod, failingHandler), Map.of(),
242+
Function.identity());
217243

218244
// Simulate incoming request that will trigger the error
219245
McpSchema.JSONRPCRequest request = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, testMethod,
@@ -232,6 +258,8 @@ void testRequestHandlerThrowsGenericException() {
232258
assertThat(response.error().data()).isNotNull();
233259
assertThat(response.error().data().toString()).contains("RuntimeException");
234260
assertThat(response.error().data().toString()).contains("Something went wrong");
261+
262+
session.close();
235263
}
236264

237265
@Test
@@ -243,8 +271,9 @@ void testRequestHandlerThrowsExceptionWithCause() {
243271
RuntimeException topException = new RuntimeException("Top level message", middleCause);
244272
McpClientSession.RequestHandler<?> failingHandler = params -> Mono.error(topException);
245273

246-
transport = new MockMcpClientTransport();
247-
session = new McpClientSession(TIMEOUT, transport, Map.of(testMethod, failingHandler), Map.of());
274+
var transport = new MockMcpClientTransport();
275+
var session = new McpClientSession(TIMEOUT, transport, Map.of(testMethod, failingHandler), Map.of(),
276+
Function.identity());
248277

249278
// Simulate incoming request that will trigger the error
250279
McpSchema.JSONRPCRequest request = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, testMethod,
@@ -267,68 +296,17 @@ void testRequestHandlerThrowsExceptionWithCause() {
267296
assertThat(dataString).contains("Middle cause message");
268297
assertThat(dataString).contains("IllegalArgumentException");
269298
assertThat(dataString).contains("Root cause message");
270-
}
271-
272-
@Test
273-
void testRequestHandlerThrowsMcpErrorWithoutJsonRpcError() {
274-
// Setup: Create a request handler that throws deprecated McpError without
275-
// JSONRPCError
276-
String testMethod = "test.deprecatedError";
277-
@SuppressWarnings("deprecation")
278-
McpError deprecatedError = new McpError("Deprecated error format");
279-
McpClientSession.RequestHandler<?> failingHandler = params -> Mono.error(deprecatedError);
280-
281-
transport = new MockMcpClientTransport();
282-
session = new McpClientSession(TIMEOUT, transport, Map.of(testMethod, failingHandler), Map.of());
283-
284-
// Simulate incoming request that will trigger the error
285-
McpSchema.JSONRPCRequest request = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, testMethod,
286-
"test-id", null);
287-
transport.simulateIncomingMessage(request);
288299

289-
// Verify: The response should create a new INTERNAL_ERROR with aggregated
290-
// messages
291-
McpSchema.JSONRPCMessage sentMessage = transport.getLastSentMessage();
292-
assertThat(sentMessage).isInstanceOf(McpSchema.JSONRPCResponse.class);
293-
McpSchema.JSONRPCResponse response = (McpSchema.JSONRPCResponse) sentMessage;
294-
assertThat(response.error()).isNotNull();
295-
assertThat(response.error().code()).isEqualTo(McpSchema.ErrorCodes.INTERNAL_ERROR);
296-
assertThat(response.error().message()).isEqualTo("Deprecated error format");
297-
// Verify data field contains aggregated exception messages
298-
assertThat(response.error().data()).isNotNull();
299-
assertThat(response.error().data().toString()).contains("McpError");
300-
assertThat(response.error().data().toString()).contains("Deprecated error format");
301-
}
302-
303-
@Test
304-
void testRequestHandlerThrowsResourceNotFoundError() {
305-
// Setup: Create a request handler that throws RESOURCE_NOT_FOUND error
306-
String testMethod = "test.resourceError";
307-
String resourceUri = "file:///missing/resource.txt";
308-
McpClientSession.RequestHandler<?> failingHandler = params -> Mono
309-
.error(McpError.RESOURCE_NOT_FOUND.apply(resourceUri));
310-
311-
transport = new MockMcpClientTransport();
312-
session = new McpClientSession(TIMEOUT, transport, Map.of(testMethod, failingHandler), Map.of());
313-
314-
// Simulate incoming request that will trigger the error
315-
McpSchema.JSONRPCRequest request = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, testMethod,
316-
"test-id", null);
317-
transport.simulateIncomingMessage(request);
318-
319-
// Verify: The response should preserve the RESOURCE_NOT_FOUND error code and
320-
// data
321-
McpSchema.JSONRPCMessage sentMessage = transport.getLastSentMessage();
322-
assertThat(sentMessage).isInstanceOf(McpSchema.JSONRPCResponse.class);
323-
McpSchema.JSONRPCResponse response = (McpSchema.JSONRPCResponse) sentMessage;
324-
assertThat(response.error()).isNotNull();
325-
assertThat(response.error().code()).isEqualTo(McpSchema.ErrorCodes.RESOURCE_NOT_FOUND);
326-
assertThat(response.error().message()).isEqualTo("Resource not found");
327-
assertThat(response.error().data()).isEqualTo(Map.of("uri", resourceUri));
300+
session.close();
328301
}
329302

330303
@Test
331304
void testGracefulShutdown() {
305+
var transport = new MockMcpClientTransport();
306+
var session = new McpClientSession(TIMEOUT, transport, Map.of(),
307+
Map.of(TEST_NOTIFICATION, params -> Mono.fromRunnable(() -> logger.info("Status update: {}", params))),
308+
Function.identity());
309+
332310
StepVerifier.create(session.closeGracefully()).verifyComplete();
333311
}
334312

0 commit comments

Comments
 (0)