WIP

mcp/src/test/java/spring/ai/mcp/McpAsyncClientTests.java

@@ -0,0 +1,151 @@
+/*
+ * 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;
+
+import java.time.Duration;
+import java.util.Map;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+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.StdioServerTransport;
+import spring.ai.mcp.spec.McpSchema.CallToolRequest;
+import spring.ai.mcp.spec.McpSchema.Resource;
+import spring.ai.mcp.spec.McpSchema.Tool;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatCode;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/**
+ * Unit tests for MCP Client Session functionality.
+ *
+ * @author Christian Tzolov
+ * @since 1.0.0
+ */
+class McpAsyncClientTests {
+
+	private McpAsyncClient mcpAsyncClient;
+
+	private StdioServerParameters stdioParams;
+
+	private static final Duration TIMEOUT = Duration.ofSeconds(10);
+
+	private static final String TEST_MESSAGE = "Hello MCP Spring AI!";
+
+	@BeforeEach
+	void setUp() {
+		stdioParams = StdioServerParameters.builder("npx")
+			.args("-y", "@modelcontextprotocol/server-everything", "dir")
+			.build();
+
+		assertThatCode(() -> {
+			mcpAsyncClient = McpClient.async(new StdioServerTransport(stdioParams), TIMEOUT, new ObjectMapper());
+			mcpAsyncClient.initialize();
+		}).doesNotThrowAnyException();
+	}
+
+	@AfterEach
+	void tearDown() {
+		// if (mcpSyncClient != null) {
+		// assertThatCode(() -> mcpSyncClient.close()).doesNotThrowAnyException();
+		// }
+	}
+
+	@Test
+	void testConstructorWithInvalidArguments() {
+		assertThatThrownBy(() -> McpClient.sync(null, TIMEOUT, new ObjectMapper()))
+			.isInstanceOf(IllegalArgumentException.class);
+
+		assertThatThrownBy(() -> McpClient.sync(new StdioServerTransport(stdioParams), null, new ObjectMapper()))
+			.isInstanceOf(IllegalArgumentException.class);
+
+		assertThatThrownBy(() -> McpClient.sync(new StdioServerTransport(stdioParams), TIMEOUT, null))
+			.isInstanceOf(IllegalArgumentException.class);
+	}
+
+	@Test
+	@Timeout(15) // Giving extra time beyond the client timeout
+	void testListTools() {
+		mcpAsyncClient.listTools(null).subscribe(result -> {
+			assertThat(result.tools()).isNotNull().isNotEmpty();
+
+			Tool firstTool = result.tools().get(0);
+			assertThat(firstTool.name()).isNotNull();
+			assertThat(firstTool.description()).isNotNull();
+		});
+	}
+
+	@Test
+	@Timeout(15)
+	void testPing() {
+		assertThatCode(() -> mcpAsyncClient.ping().block()).doesNotThrowAnyException();
+	}
+
+	@Test
+	@Timeout(15)
+	void testCallTool() {
+		CallToolRequest callToolRequest = new CallToolRequest("echo", Map.of("message", TEST_MESSAGE));
+
+		mcpAsyncClient.callTool(callToolRequest).subscribe(callToolResult -> {
+			assertThat(callToolResult).isNotNull().satisfies(result -> {
+				assertThat(result.content()).isNotNull();
+				assertThat(result.isError()).isNull();
+			});
+		});
+	}
+
+	@Test
+	@Timeout(15)
+	void testCallToolWithInvalidTool() {
+		CallToolRequest invalidRequest = new CallToolRequest("nonexistent_tool", Map.of("message", TEST_MESSAGE));
+
+		assertThatThrownBy(() -> mcpAsyncClient.callTool(invalidRequest).block()).isInstanceOf(Exception.class);
+	}
+
+	@Test
+	@Timeout(15)
+	void testRootsListChanged() {
+		assertThatCode(() -> mcpAsyncClient.sendRootsListChanged().block()).doesNotThrowAnyException();
+	}
+
+	@Test
+	@Timeout(15)
+	void testListResources() {
+		mcpAsyncClient.listResources(null).subscribe(resources -> {
+			assertThat(resources).isNotNull().satisfies(result -> {
+				assertThat(result.resources()).isNotNull();
+
+				if (!result.resources().isEmpty()) {
+					Resource firstResource = result.resources().get(0);
+					assertThat(firstResource.uri()).isNotNull();
+					assertThat(firstResource.name()).isNotNull();
+				}
+			});
+		});
+	}
+
+	@Test
+	void testMcpAsyncClientState() {
+		assertThat(mcpAsyncClient).isNotNull();
+	}
+
+}

mcp/src/test/java/spring/ai/mcp/McpSyncClientTests.java

@@ -0,0 +1,155 @@
+/*
+ * 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;
+
+import java.time.Duration;
+import java.util.Map;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.Timeout;
+
+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.StdioServerTransport;
+import spring.ai.mcp.spec.McpSchema.CallToolRequest;
+import spring.ai.mcp.spec.McpSchema.CallToolResult;
+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;
+
+/**
+ * Unit tests for MCP Client Session functionality.
+ *
+ * @author Christian Tzolov
+ * @since 1.0.0
+ */
+class McpSyncClientTests {
+
+	private McpSyncClient mcpSyncClient;
+
+	private StdioServerParameters stdioParams;
+
+	private static final Duration TIMEOUT = Duration.ofSeconds(10);
+
+	private static final String TEST_MESSAGE = "Hello MCP Spring AI!";
+
+	@BeforeEach
+	void setUp() {
+		stdioParams = StdioServerParameters.builder("npx")
+			.args("-y", "@modelcontextprotocol/server-everything", "dir")
+			.build();
+
+		assertThatCode(() -> {
+			mcpSyncClient = McpClient.sync(new StdioServerTransport(stdioParams), TIMEOUT, new ObjectMapper());
+			mcpSyncClient.initialize();
+		}).doesNotThrowAnyException();
+	}
+
+	@AfterEach
+	void tearDown() {
+		if (mcpSyncClient != null) {
+			assertThatCode(() -> mcpSyncClient.close()).doesNotThrowAnyException();
+		}
+	}
+
+	@Test
+	void testConstructorWithInvalidArguments() {
+		assertThatThrownBy(() -> McpClient.sync(null, TIMEOUT, new ObjectMapper()))
+			.isInstanceOf(IllegalArgumentException.class);
+
+		assertThatThrownBy(() -> McpClient.sync(new StdioServerTransport(stdioParams), null, new ObjectMapper()))
+			.isInstanceOf(IllegalArgumentException.class);
+
+		assertThatThrownBy(() -> McpClient.sync(new StdioServerTransport(stdioParams), TIMEOUT, null))
+			.isInstanceOf(IllegalArgumentException.class);
+	}
+
+	@Test
+	@Timeout(15) // Giving extra time beyond the client timeout
+	void testListTools() {
+		ListToolsResult tools = mcpSyncClient.listTools(null);
+
+		assertThat(tools).isNotNull().satisfies(result -> {
+			assertThat(result.tools()).isNotNull().isNotEmpty();
+
+			Tool firstTool = result.tools().get(0);
+			assertThat(firstTool.name()).isNotNull();
+			assertThat(firstTool.description()).isNotNull();
+		});
+	}
+
+	@Test
+	@Timeout(15)
+	void testPing() {
+		assertThatCode(() -> mcpSyncClient.ping()).doesNotThrowAnyException();
+	}
+
+	@Test
+	@Timeout(15)
+	void testCallTool() {
+		CallToolRequest callToolRequest = new CallToolRequest("echo", Map.of("message", TEST_MESSAGE));
+
+		CallToolResult callToolResult = mcpSyncClient.callTool(callToolRequest);
+
+		assertThat(callToolResult).isNotNull().satisfies(result -> {
+			assertThat(result.content()).isNotNull();
+			assertThat(result.isError()).isNull();
+		});
+	}
+
+	@Test
+	@Timeout(15)
+	void testCallToolWithInvalidTool() {
+		CallToolRequest invalidRequest = new CallToolRequest("nonexistent_tool", Map.of("message", TEST_MESSAGE));
+
+		assertThatThrownBy(() -> mcpSyncClient.callTool(invalidRequest)).isInstanceOf(Exception.class);
+	}
+
+	@Test
+	@Timeout(15)
+	void testRootsListChanged() {
+		assertThatCode(() -> mcpSyncClient.sendRootsListChanged()).doesNotThrowAnyException();
+	}
+
+	@Test
+	@Timeout(15)
+	void testListResources() {
+		ListResourcesResult resources = mcpSyncClient.listResources(null);
+
+		assertThat(resources).isNotNull().satisfies(result -> {
+			assertThat(result.resources()).isNotNull();
+
+			if (!result.resources().isEmpty()) {
+				Resource firstResource = result.resources().get(0);
+				assertThat(firstResource.uri()).isNotNull();
+				assertThat(firstResource.name()).isNotNull();
+			}
+		});
+	}
+
+	@Test
+	void testClientSessionState() {
+		assertThat(mcpSyncClient).isNotNull();
+	}
+
+}

mcp/src/test/java/spring/ai/mcp/attic/ClientSessionTests2.java

@@ -0,0 +1,80 @@
+/*
+* 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.attic;
+
+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.StdioServerTransport;
+import spring.ai.mcp.spec.McpSchema.ListResourcesResult;
+import spring.ai.mcp.spec.McpSchema.ListToolsResult;
+import spring.ai.mcp.spec.McpSchema.Resource;
+
+/**
+ * @author Christian Tzolov
+ * @since 1.0.0
+ */
+public class ClientSessionTests2 {
+
+	public static void main(String[] args) {
+
+		var stdioParams = StdioServerParameters.builder("uv")
+			.args("--directory", "dir", "run", "mcp-server-sqlite",
+					"--db-path", "~/test.db")
+			.build();
+
+		try (McpSyncClient clientSession =
+				     McpClient.sync(new StdioServerTransport(stdioParams),
+						     Duration.ofSeconds(10), new ObjectMapper())) {
+
+			clientSession.initialize();
+
+			ListToolsResult tools = clientSession.listTools(null);
+			System.out.println("Tools: " + tools);
+
+			clientSession.ping();
+
+			// CallToolRequest callToolRequest = new CallToolRequest("echo",
+			// Map.of("message", "Hello MCP Spring AI!"));
+			// CallToolResult callToolResult =
+			// clientSession.callTool(callToolRequest);
+			// System.out.println("Call Tool Result: " + callToolResult);
+			// //
+			// clientSession.sendRootsListChanged();
+
+			// Resources
+			ListResourcesResult resources = clientSession.listResources(null);
+			System.out.println("Resources Size: " + resources.resources().size());
+			System.out.println("Resources: " + resources);
+			for (Resource resource : resources.resources()) {
+				System.out.println(clientSession.readResource(resource));
+
+			}
+
+			// var resourceTemplate = clientSession.listResourceTemplates(null);
+			// System.out.println("Resource Templates: " + resourceTemplate);
+
+			// stdioClient.awaitForExit();
+		}
+		catch (Exception e) {
+			e.printStackTrace();
+		}
+	}
+
+}

mcp/src/test/java/spring/ai/mcp/attic/ClientSessionTests3.java

@@ -0,0 +1,75 @@
+/*
+* 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.attic;
+
+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.StdioServerTransport;
+import spring.ai.mcp.spec.McpSchema.ListResourcesResult;
+import spring.ai.mcp.spec.McpSchema.ListToolsResult;
+import spring.ai.mcp.spec.McpSchema.Resource;
+
+/**
+ * @author Christian Tzolov
+ * @since 1.0.0
+ */
+public class ClientSessionTests3 {
+
+	public static void main(String[] args) {
+
+		var stdioParams = StdioServerParameters.builder("uv")
+			.args("--directory", "dir", "run", "mcp-server-sqlite",
+					"--db-path", "~/test.db")
+			.build();
+
+
+		McpSyncClient clientSession = null;
+		try {
+
+			clientSession = McpClient.sync(new StdioServerTransport(stdioParams),
+				Duration.ofSeconds(10), new ObjectMapper());
+
+			clientSession.initialize();
+
+			ListToolsResult tools = clientSession.listTools(null);
+			System.out.println("Tools: " + tools);
+
+			clientSession.ping();
+
+			// Resources
+			ListResourcesResult resources = clientSession.listResources(null);
+			System.out.println("Resources Size: " + resources.resources().size());
+			System.out.println("Resources: " + resources);
+			for (Resource resource : resources.resources()) {
+				System.out.println(clientSession.readResource(resource));
+
+			}
+		}
+		catch (Exception e) {
+			e.printStackTrace();
+		}
+		finally {
+			if (clientSession != null) {
+				clientSession.close();
+			}
+		}
+	}
+
+}

mcp/src/test/java/spring/ai/mcp/spec/McpSessionJSONRPCTest.java

@@ -0,0 +1,158 @@
+/* 
+* 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();
+    }
+}

mcp/src/test/java/spring/ai/mcp/spec/McpSessionTest.java

@@ -0,0 +1,93 @@
+/*
+* 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();
+	}
+
+}

mcp/src/test/java/spring/ai/mcp/spec/McpTransportTest.java

@@ -0,0 +1,128 @@
+/* 
+* 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();
+    }
+}