Align abstract tests in mcp-test with mcp-core (#70)

Signed-off-by: Dariusz Jędrzejczyk <dariusz.jedrzejczyk@broadcom.com>

mcp-test/src/main/java/org/springframework/ai/mcp/client/AbstractMcpAsyncClientTests.java

@@ -0,0 +1,375 @@
+/*
+ * 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 org.springframework.ai.mcp.client;
+
+import java.time.Duration;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.Function;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Disabled;
+import org.junit.jupiter.api.Test;
+import reactor.core.publisher.Mono;
+import reactor.test.StepVerifier;
+
+import org.springframework.ai.mcp.spec.ClientMcpTransport;
+import org.springframework.ai.mcp.spec.McpSchema;
+import org.springframework.ai.mcp.spec.McpSchema.CallToolRequest;
+import org.springframework.ai.mcp.spec.McpSchema.ClientCapabilities;
+import org.springframework.ai.mcp.spec.McpSchema.CreateMessageRequest;
+import org.springframework.ai.mcp.spec.McpSchema.CreateMessageResult;
+import org.springframework.ai.mcp.spec.McpSchema.GetPromptRequest;
+import org.springframework.ai.mcp.spec.McpSchema.Prompt;
+import org.springframework.ai.mcp.spec.McpSchema.Resource;
+import org.springframework.ai.mcp.spec.McpSchema.Root;
+import org.springframework.ai.mcp.spec.McpSchema.SubscribeRequest;
+import org.springframework.ai.mcp.spec.McpSchema.Tool;
+import org.springframework.ai.mcp.spec.McpSchema.UnsubscribeRequest;
+import org.springframework.ai.mcp.spec.McpTransport;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatCode;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/**
+ * Test suite for the {@link McpAsyncClient} that can be used with different
+ * {@link McpTransport} implementations.
+ *
+ * @author Christian Tzolov
+ * @author Dariusz Jędrzejczyk
+ */
+public abstract class AbstractMcpAsyncClientTests {
+
+	private McpAsyncClient mcpAsyncClient;
+
+	protected ClientMcpTransport mcpTransport;
+
+	private static final Duration TIMEOUT = Duration.ofSeconds(20);
+
+	private static final String ECHO_TEST_MESSAGE = "Hello MCP Spring AI!";
+
+	abstract protected ClientMcpTransport createMcpTransport();
+
+	protected void onStart() {
+	}
+
+	protected void onClose() {
+	}
+
+	@BeforeEach
+	void setUp() {
+		onStart();
+		this.mcpTransport = createMcpTransport();
+
+		assertThatCode(() -> {
+			mcpAsyncClient = McpClient.async(mcpTransport)
+				.requestTimeout(TIMEOUT)
+				.capabilities(ClientCapabilities.builder().roots(true).build())
+				.build();
+			mcpAsyncClient.initialize().block(Duration.ofSeconds(10));
+		}).doesNotThrowAnyException();
+	}
+
+	@AfterEach
+	void tearDown() {
+		if (mcpAsyncClient != null) {
+			assertThatCode(() -> mcpAsyncClient.closeGracefully().block(Duration.ofSeconds(10)))
+				.doesNotThrowAnyException();
+		}
+		onClose();
+	}
+
+	@Test
+	void testConstructorWithInvalidArguments() {
+		assertThatThrownBy(() -> McpClient.sync(null).build()).isInstanceOf(IllegalArgumentException.class)
+			.hasMessage("Transport must not be null");
+
+		assertThatThrownBy(() -> McpClient.sync(mcpTransport).requestTimeout(null).build())
+			.isInstanceOf(IllegalArgumentException.class)
+			.hasMessage("Request timeout must not be null");
+	}
+
+	@Test
+	void testListTools() {
+		StepVerifier.create(mcpAsyncClient.listTools(null)).consumeNextWith(result -> {
+			assertThat(result.tools()).isNotNull().isNotEmpty();
+
+			Tool firstTool = result.tools().get(0);
+			assertThat(firstTool.name()).isNotNull();
+			assertThat(firstTool.description()).isNotNull();
+		}).verifyComplete();
+	}
+
+	@Test
+	void testPing() {
+		assertThatCode(() -> mcpAsyncClient.ping().block()).doesNotThrowAnyException();
+	}
+
+	@Test
+	void testCallTool() {
+		CallToolRequest callToolRequest = new CallToolRequest("echo", Map.of("message", ECHO_TEST_MESSAGE));
+
+		StepVerifier.create(mcpAsyncClient.callTool(callToolRequest)).consumeNextWith(callToolResult -> {
+			assertThat(callToolResult).isNotNull().satisfies(result -> {
+				assertThat(result.content()).isNotNull();
+				assertThat(result.isError()).isNull();
+			});
+		}).verifyComplete();
+	}
+
+	@Test
+	void testCallToolWithInvalidTool() {
+		CallToolRequest invalidRequest = new CallToolRequest("nonexistent_tool", Map.of("message", ECHO_TEST_MESSAGE));
+
+		assertThatThrownBy(() -> mcpAsyncClient.callTool(invalidRequest).block()).isInstanceOf(Exception.class);
+	}
+
+	@Test
+	void testListResources() {
+		StepVerifier.create(mcpAsyncClient.listResources(null)).consumeNextWith(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();
+				}
+			});
+		}).verifyComplete();
+	}
+
+	@Test
+	void testMcpAsyncClientState() {
+		assertThat(mcpAsyncClient).isNotNull();
+	}
+
+	@Test
+	void testListPrompts() {
+		StepVerifier.create(mcpAsyncClient.listPrompts(null)).consumeNextWith(prompts -> {
+			assertThat(prompts).isNotNull().satisfies(result -> {
+				assertThat(result.prompts()).isNotNull();
+
+				if (!result.prompts().isEmpty()) {
+					Prompt firstPrompt = result.prompts().get(0);
+					assertThat(firstPrompt.name()).isNotNull();
+					assertThat(firstPrompt.description()).isNotNull();
+				}
+			});
+		}).verifyComplete();
+	}
+
+	@Test
+	void testGetPrompt() {
+		StepVerifier.create(mcpAsyncClient.getPrompt(new GetPromptRequest("simple_prompt", Map.of())))
+			.consumeNextWith(prompt -> {
+				assertThat(prompt).isNotNull().satisfies(result -> {
+					assertThat(result.messages()).isNotEmpty();
+					assertThat(result.messages()).hasSize(1);
+				});
+			})
+			.verifyComplete();
+	}
+
+	@Test
+	void testRootsListChanged() {
+		assertThatCode(() -> mcpAsyncClient.rootsListChangedNotification().block()).doesNotThrowAnyException();
+	}
+
+	@Test
+	void testInitializeWithRootsListProviders() {
+		var transport = createMcpTransport();
+
+		var client = McpClient.async(transport)
+			.requestTimeout(TIMEOUT)
+			.roots(new Root("file:///test/path", "test-root"))
+			.build();
+
+		assertThatCode(() -> client.initialize().block(Duration.ofSeconds(10))).doesNotThrowAnyException();
+
+		assertThatCode(() -> client.closeGracefully().block(Duration.ofSeconds(10))).doesNotThrowAnyException();
+	}
+
+	@Test
+	void testAddRoot() {
+		Root newRoot = new Root("file:///new/test/path", "new-test-root");
+		assertThatCode(() -> mcpAsyncClient.addRoot(newRoot).block()).doesNotThrowAnyException();
+	}
+
+	@Test
+	void testAddRootWithNullValue() {
+		assertThatThrownBy(() -> mcpAsyncClient.addRoot(null).block()).hasMessageContaining("Root must not be null");
+	}
+
+	@Test
+	void testRemoveRoot() {
+		Root root = new Root("file:///test/path/to/remove", "root-to-remove");
+		assertThatCode(() -> {
+			mcpAsyncClient.addRoot(root).block();
+			mcpAsyncClient.removeRoot(root.uri()).block();
+		}).doesNotThrowAnyException();
+	}
+
+	@Test
+	void testRemoveNonExistentRoot() {
+		assertThatThrownBy(() -> mcpAsyncClient.removeRoot("nonexistent-uri").block())
+			.hasMessageContaining("Root with uri 'nonexistent-uri' not found");
+	}
+
+	@Test
+	@Disabled
+	void testReadResource() {
+		StepVerifier.create(mcpAsyncClient.listResources()).consumeNextWith(resources -> {
+			if (!resources.resources().isEmpty()) {
+				Resource firstResource = resources.resources().get(0);
+				StepVerifier.create(mcpAsyncClient.readResource(firstResource)).consumeNextWith(result -> {
+					assertThat(result).isNotNull();
+					assertThat(result.contents()).isNotNull();
+				}).verifyComplete();
+			}
+		}).verifyComplete();
+	}
+
+	@Test
+	void testListResourceTemplates() {
+		StepVerifier.create(mcpAsyncClient.listResourceTemplates()).consumeNextWith(result -> {
+			assertThat(result).isNotNull();
+			assertThat(result.resourceTemplates()).isNotNull();
+		}).verifyComplete();
+	}
+
+	// @Test
+	void testResourceSubscription() {
+		StepVerifier.create(mcpAsyncClient.listResources()).consumeNextWith(resources -> {
+			if (!resources.resources().isEmpty()) {
+				Resource firstResource = resources.resources().get(0);
+
+				// Test subscribe
+				StepVerifier.create(mcpAsyncClient.subscribeResource(new SubscribeRequest(firstResource.uri())))
+					.verifyComplete();
+
+				// Test unsubscribe
+				StepVerifier.create(mcpAsyncClient.unsubscribeResource(new UnsubscribeRequest(firstResource.uri())))
+					.verifyComplete();
+			}
+		}).verifyComplete();
+	}
+
+	@Test
+	void testNotificationHandlers() {
+		AtomicBoolean toolsNotificationReceived = new AtomicBoolean(false);
+		AtomicBoolean resourcesNotificationReceived = new AtomicBoolean(false);
+		AtomicBoolean promptsNotificationReceived = new AtomicBoolean(false);
+
+		var transport = createMcpTransport();
+		var client = McpClient.async(transport)
+			.requestTimeout(TIMEOUT)
+			.toolsChangeConsumer(tools -> Mono.fromRunnable(() -> toolsNotificationReceived.set(true)))
+			.resourcesChangeConsumer(resources -> Mono.fromRunnable(() -> resourcesNotificationReceived.set(true)))
+			.promptsChangeConsumer(prompts -> Mono.fromRunnable(() -> promptsNotificationReceived.set(true)))
+			.build();
+
+		assertThatCode(() -> {
+			client.initialize().block();
+			// Trigger notifications
+			client.sendResourcesListChanged().block();
+			client.promptListChangedNotification().block();
+			client.closeGracefully().block();
+		}).doesNotThrowAnyException();
+	}
+
+	@Test
+	void testInitializeWithSamplingCapability() {
+		var transport = createMcpTransport();
+
+		var capabilities = ClientCapabilities.builder().sampling().build();
+
+		var client = McpClient.async(transport)
+			.requestTimeout(TIMEOUT)
+			.capabilities(capabilities)
+			.sampling(request -> Mono.just(CreateMessageResult.builder().message("test").model("test-model").build()))
+			.build();
+
+		assertThatCode(() -> {
+			client.initialize().block(Duration.ofSeconds(10));
+			client.closeGracefully().block(Duration.ofSeconds(10));
+		}).doesNotThrowAnyException();
+	}
+
+	@Test
+	void testInitializeWithAllCapabilities() {
+		var transport = createMcpTransport();
+
+		var capabilities = ClientCapabilities.builder()
+			.experimental(Map.of("feature", "test"))
+			.roots(true)
+			.sampling()
+			.build();
+
+		Function<CreateMessageRequest, Mono<CreateMessageResult>> samplingHandler = request -> Mono
+			.just(CreateMessageResult.builder().message("test").model("test-model").build());
+		var client = McpClient.async(transport)
+			.requestTimeout(TIMEOUT)
+			.capabilities(capabilities)
+			.sampling(samplingHandler)
+			.build();
+
+		assertThatCode(() -> {
+			var result = client.initialize().block(Duration.ofSeconds(10));
+			assertThat(result).isNotNull();
+			assertThat(result.capabilities()).isNotNull();
+			client.closeGracefully().block(Duration.ofSeconds(10));
+		}).doesNotThrowAnyException();
+	}
+
+	// ---------------------------------------
+	// Logging Tests
+	// ---------------------------------------
+
+	@Test
+	void testLoggingLevels() {
+		// Test all logging levels
+		for (McpSchema.LoggingLevel level : McpSchema.LoggingLevel.values()) {
+			StepVerifier.create(mcpAsyncClient.setLoggingLevel(level)).verifyComplete();
+		}
+	}
+
+	@Test
+	void testLoggingConsumer() {
+		AtomicBoolean logReceived = new AtomicBoolean(false);
+		var transport = createMcpTransport();
+
+		var client = McpClient.async(transport)
+			.requestTimeout(TIMEOUT)
+			.loggingConsumer(notification -> Mono.fromRunnable(() -> logReceived.set(true)))
+			.build();
+
+		assertThatCode(() -> {
+			client.initialize().block(Duration.ofSeconds(10));
+			client.closeGracefully().block(Duration.ofSeconds(10));
+		}).doesNotThrowAnyException();
+	}
+
+	@Test
+	void testLoggingWithNullNotification() {
+		assertThatThrownBy(() -> mcpAsyncClient.setLoggingLevel(null).block())
+			.hasMessageContaining("Logging level must not be null");
+	}
+
+}

mcp-test/src/main/java/org/springframework/ai/mcp/client/AbstractMcpSyncClientTests.java

@@ -0,0 +1,315 @@
+/*
+ * 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 org.springframework.ai.mcp.client;
+
+import java.time.Duration;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import org.springframework.ai.mcp.spec.ClientMcpTransport;
+import org.springframework.ai.mcp.spec.McpSchema;
+import org.springframework.ai.mcp.spec.McpSchema.CallToolRequest;
+import org.springframework.ai.mcp.spec.McpSchema.CallToolResult;
+import org.springframework.ai.mcp.spec.McpSchema.ClientCapabilities;
+import org.springframework.ai.mcp.spec.McpSchema.ListResourceTemplatesResult;
+import org.springframework.ai.mcp.spec.McpSchema.ListResourcesResult;
+import org.springframework.ai.mcp.spec.McpSchema.ListToolsResult;
+import org.springframework.ai.mcp.spec.McpSchema.ReadResourceResult;
+import org.springframework.ai.mcp.spec.McpSchema.Resource;
+import org.springframework.ai.mcp.spec.McpSchema.Root;
+import org.springframework.ai.mcp.spec.McpSchema.SubscribeRequest;
+import org.springframework.ai.mcp.spec.McpSchema.TextContent;
+import org.springframework.ai.mcp.spec.McpSchema.Tool;
+import org.springframework.ai.mcp.spec.McpSchema.UnsubscribeRequest;
+
+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
+ * @author Dariusz Jędrzejczyk
+ */
+public abstract class AbstractMcpSyncClientTests {
+
+	private McpSyncClient mcpSyncClient;
+
+	private static final Duration TIMEOUT = Duration.ofSeconds(10);
+
+	private static final String TEST_MESSAGE = "Hello MCP Spring AI!";
+
+	protected ClientMcpTransport mcpTransport;
+
+	abstract protected ClientMcpTransport createMcpTransport();
+
+	abstract protected void onStart();
+
+	abstract protected void onClose();
+
+	@BeforeEach
+	void setUp() {
+		onStart();
+		this.mcpTransport = createMcpTransport();
+
+		assertThatCode(() -> {
+			mcpSyncClient = McpClient.sync(mcpTransport)
+				.requestTimeout(TIMEOUT)
+				.capabilities(ClientCapabilities.builder().roots(true).build())
+				.build();
+			mcpSyncClient.initialize();
+		}).doesNotThrowAnyException();
+	}
+
+	@AfterEach
+	void tearDown() {
+		if (mcpSyncClient != null) {
+			assertThatCode(() -> mcpSyncClient.close()).doesNotThrowAnyException();
+		}
+		onClose();
+	}
+
+	@Test
+	void testConstructorWithInvalidArguments() {
+		assertThatThrownBy(() -> McpClient.sync(null).build()).isInstanceOf(IllegalArgumentException.class)
+			.hasMessage("Transport must not be null");
+
+		assertThatThrownBy(() -> McpClient.sync(mcpTransport).requestTimeout(null).build())
+			.isInstanceOf(IllegalArgumentException.class)
+			.hasMessage("Request timeout must not be null");
+	}
+
+	@Test
+	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
+	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");
+		});
+	}
+
+	@Test
+	void testPing() {
+		assertThatCode(() -> mcpSyncClient.ping()).doesNotThrowAnyException();
+	}
+
+	@Test
+	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
+	void testCallToolWithInvalidTool() {
+		CallToolRequest invalidRequest = new CallToolRequest("nonexistent_tool", Map.of("message", TEST_MESSAGE));
+
+		assertThatThrownBy(() -> mcpSyncClient.callTool(invalidRequest)).isInstanceOf(Exception.class);
+	}
+
+	@Test
+	void testRootsListChanged() {
+		assertThatCode(() -> mcpSyncClient.rootsListChangedNotification()).doesNotThrowAnyException();
+	}
+
+	@Test
+	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();
+	}
+
+	@Test
+	void testInitializeWithRootsListProviders() {
+		var transport = createMcpTransport();
+
+		var client = McpClient.sync(transport)
+			.requestTimeout(TIMEOUT)
+			.roots(new Root("file:///test/path", "test-root"))
+			.build();
+
+		assertThatCode(() -> {
+			client.initialize();
+			client.close();
+		}).doesNotThrowAnyException();
+	}
+
+	@Test
+	void testAddRoot() {
+		Root newRoot = new Root("file:///new/test/path", "new-test-root");
+		assertThatCode(() -> mcpSyncClient.addRoot(newRoot)).doesNotThrowAnyException();
+	}
+
+	@Test
+	void testAddRootWithNullValue() {
+		assertThatThrownBy(() -> mcpSyncClient.addRoot(null)).hasMessageContaining("Root must not be null");
+	}
+
+	@Test
+	void testRemoveRoot() {
+		Root root = new Root("file:///test/path/to/remove", "root-to-remove");
+		assertThatCode(() -> {
+			mcpSyncClient.addRoot(root);
+			mcpSyncClient.removeRoot(root.uri());
+		}).doesNotThrowAnyException();
+	}
+
+	@Test
+	void testRemoveNonExistentRoot() {
+		assertThatThrownBy(() -> mcpSyncClient.removeRoot("nonexistent-uri"))
+			.hasMessageContaining("Root with uri 'nonexistent-uri' not found");
+	}
+
+	@Test
+	void testReadResource() {
+		ListResourcesResult resources = mcpSyncClient.listResources(null);
+
+		if (!resources.resources().isEmpty()) {
+			Resource firstResource = resources.resources().get(0);
+			ReadResourceResult result = mcpSyncClient.readResource(firstResource);
+
+			assertThat(result).isNotNull();
+			assertThat(result.contents()).isNotNull();
+		}
+	}
+
+	@Test
+	void testListResourceTemplates() {
+		ListResourceTemplatesResult result = mcpSyncClient.listResourceTemplates(null);
+
+		assertThat(result).isNotNull();
+		assertThat(result.resourceTemplates()).isNotNull();
+	}
+
+	// @Test
+	void testResourceSubscription() {
+		ListResourcesResult resources = mcpSyncClient.listResources(null);
+
+		if (!resources.resources().isEmpty()) {
+			Resource firstResource = resources.resources().get(0);
+
+			// Test subscribe
+			assertThatCode(() -> mcpSyncClient.subscribeResource(new SubscribeRequest(firstResource.uri())))
+				.doesNotThrowAnyException();
+
+			// Test unsubscribe
+			assertThatCode(() -> mcpSyncClient.unsubscribeResource(new UnsubscribeRequest(firstResource.uri())))
+				.doesNotThrowAnyException();
+		}
+	}
+
+	@Test
+	void testNotificationHandlers() {
+		AtomicBoolean toolsNotificationReceived = new AtomicBoolean(false);
+		AtomicBoolean resourcesNotificationReceived = new AtomicBoolean(false);
+		AtomicBoolean promptsNotificationReceived = new AtomicBoolean(false);
+
+		var transport = createMcpTransport();
+		var client = McpClient.sync(transport)
+			.requestTimeout(TIMEOUT)
+			.toolsChangeConsumer(tools -> toolsNotificationReceived.set(true))
+			.resourcesChangeConsumer(resources -> resourcesNotificationReceived.set(true))
+			.promptsChangeConsumer(prompts -> promptsNotificationReceived.set(true))
+			.build();
+
+		assertThatCode(() -> {
+			client.initialize();
+			// Trigger notifications
+			client.sendResourcesListChanged();
+			client.promptListChangedNotification();
+			client.close();
+		}).doesNotThrowAnyException();
+	}
+
+	// ---------------------------------------
+	// Logging Tests
+	// ---------------------------------------
+
+	@Test
+	void testLoggingLevels() {
+		// Test all logging levels
+		for (McpSchema.LoggingLevel level : McpSchema.LoggingLevel.values()) {
+			assertThatCode(() -> mcpSyncClient.setLoggingLevel(level)).doesNotThrowAnyException();
+		}
+	}
+
+	@Test
+	void testLoggingConsumer() {
+		AtomicBoolean logReceived = new AtomicBoolean(false);
+		var transport = createMcpTransport();
+
+		var client = McpClient.sync(transport)
+			.requestTimeout(TIMEOUT)
+			.loggingConsumer(notification -> logReceived.set(true))
+			.build();
+
+		assertThatCode(() -> {
+			client.initialize();
+			client.close();
+		}).doesNotThrowAnyException();
+	}
+
+	@Test
+	void testLoggingWithNullNotification() {
+		assertThatThrownBy(() -> mcpSyncClient.setLoggingLevel(null))
+			.hasMessageContaining("Logging level must not be null");
+	}
+
+}

mcp-test/src/main/java/org/springframework/ai/mcp/server/AbstractMcpAsyncServerTests.java

@@ -0,0 +1,476 @@
+/*
+ * 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 org.springframework.ai.mcp.server;
+
+import java.time.Duration;
+import java.util.List;
+
+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.test.StepVerifier;
+
+import org.springframework.ai.mcp.spec.McpError;
+import org.springframework.ai.mcp.spec.McpSchema;
+import org.springframework.ai.mcp.spec.McpSchema.CallToolResult;
+import org.springframework.ai.mcp.spec.McpSchema.GetPromptResult;
+import org.springframework.ai.mcp.spec.McpSchema.Prompt;
+import org.springframework.ai.mcp.spec.McpSchema.PromptMessage;
+import org.springframework.ai.mcp.spec.McpSchema.ReadResourceResult;
+import org.springframework.ai.mcp.spec.McpSchema.Resource;
+import org.springframework.ai.mcp.spec.McpSchema.ServerCapabilities;
+import org.springframework.ai.mcp.spec.McpSchema.Tool;
+import org.springframework.ai.mcp.spec.McpTransport;
+import org.springframework.ai.mcp.spec.ServerMcpTransport;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatCode;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/**
+ * Test suite for the {@link McpAsyncServer} that can be used with different
+ * {@link McpTransport} implementations.
+ *
+ * @author Christian Tzolov
+ */
+public abstract class AbstractMcpAsyncServerTests {
+
+	private static final String TEST_TOOL_NAME = "test-tool";
+
+	private static final String TEST_RESOURCE_URI = "test://resource";
+
+	private static final String TEST_PROMPT_NAME = "test-prompt";
+
+	abstract protected ServerMcpTransport createMcpTransport();
+
+	protected void onStart() {
+	}
+
+	protected void onClose() {
+	}
+
+	@BeforeEach
+	void setUp() {
+	}
+
+	@AfterEach
+	void tearDown() {
+		onClose();
+	}
+
+	// ---------------------------------------
+	// Server Lifecycle Tests
+	// ---------------------------------------
+
+	@Test
+	void testConstructorWithInvalidArguments() {
+		assertThatThrownBy(() -> McpServer.async(null)).isInstanceOf(IllegalArgumentException.class)
+			.hasMessage("Transport must not be null");
+
+		assertThatThrownBy(() -> McpServer.async(createMcpTransport()).serverInfo((McpSchema.Implementation) null))
+			.isInstanceOf(IllegalArgumentException.class)
+			.hasMessage("Server info must not be null");
+	}
+
+	@Test
+	void testGracefulShutdown() {
+		var mcpAsyncServer = McpServer.async(createMcpTransport()).serverInfo("test-server", "1.0.0").build();
+
+		StepVerifier.create(mcpAsyncServer.closeGracefully()).verifyComplete();
+	}
+
+	@Test
+	void testImmediateClose() {
+		var mcpAsyncServer = McpServer.async(createMcpTransport()).serverInfo("test-server", "1.0.0").build();
+
+		assertThatCode(() -> mcpAsyncServer.close()).doesNotThrowAnyException();
+	}
+
+	// ---------------------------------------
+	// Tools Tests
+	// ---------------------------------------
+	String emptyJsonSchema = """
+			{
+				"$schema": "http://json-schema.org/draft-07/schema#",
+				"type": "object",
+				"properties": {}
+			}
+			""";
+
+	@Test
+	void testAddTool() {
+		Tool newTool = new McpSchema.Tool("new-tool", "New test tool", emptyJsonSchema);
+		var mcpAsyncServer = McpServer.async(createMcpTransport())
+			.serverInfo("test-server", "1.0.0")
+			.capabilities(ServerCapabilities.builder().tools(true).build())
+			.build();
+
+		StepVerifier.create(mcpAsyncServer.addTool(new McpServerFeatures.AsyncToolRegistration(newTool,
+				args -> Mono.just(new CallToolResult(List.of(), false)))))
+			.verifyComplete();
+
+		assertThatCode(() -> mcpAsyncServer.closeGracefully().block(Duration.ofSeconds(10))).doesNotThrowAnyException();
+	}
+
+	@Test
+	void testAddDuplicateTool() {
+		Tool duplicateTool = new McpSchema.Tool(TEST_TOOL_NAME, "Duplicate tool", emptyJsonSchema);
+
+		var mcpAsyncServer = McpServer.async(createMcpTransport())
+			.serverInfo("test-server", "1.0.0")
+			.capabilities(ServerCapabilities.builder().tools(true).build())
+			.tool(duplicateTool, args -> Mono.just(new CallToolResult(List.of(), false)))
+			.build();
+
+		StepVerifier.create(mcpAsyncServer.addTool(new McpServerFeatures.AsyncToolRegistration(duplicateTool,
+				args -> Mono.just(new CallToolResult(List.of(), false)))))
+			.verifyErrorSatisfies(error -> {
+				assertThat(error).isInstanceOf(McpError.class)
+					.hasMessage("Tool with name '" + TEST_TOOL_NAME + "' already exists");
+			});
+
+		assertThatCode(() -> mcpAsyncServer.closeGracefully().block(Duration.ofSeconds(10))).doesNotThrowAnyException();
+	}
+
+	@Test
+	void testRemoveTool() {
+		Tool too = new McpSchema.Tool(TEST_TOOL_NAME, "Duplicate tool", emptyJsonSchema);
+
+		var mcpAsyncServer = McpServer.async(createMcpTransport())
+			.serverInfo("test-server", "1.0.0")
+			.capabilities(ServerCapabilities.builder().tools(true).build())
+			.tool(too, args -> Mono.just(new CallToolResult(List.of(), false)))
+			.build();
+
+		StepVerifier.create(mcpAsyncServer.removeTool(TEST_TOOL_NAME)).verifyComplete();
+
+		assertThatCode(() -> mcpAsyncServer.closeGracefully().block(Duration.ofSeconds(10))).doesNotThrowAnyException();
+	}
+
+	@Test
+	void testRemoveNonexistentTool() {
+		var mcpAsyncServer = McpServer.async(createMcpTransport())
+			.serverInfo("test-server", "1.0.0")
+			.capabilities(ServerCapabilities.builder().tools(true).build())
+			.build();
+
+		StepVerifier.create(mcpAsyncServer.removeTool("nonexistent-tool")).verifyErrorSatisfies(error -> {
+			assertThat(error).isInstanceOf(McpError.class).hasMessage("Tool with name 'nonexistent-tool' not found");
+		});
+
+		assertThatCode(() -> mcpAsyncServer.closeGracefully().block(Duration.ofSeconds(10))).doesNotThrowAnyException();
+	}
+
+	@Test
+	void testNotifyToolsListChanged() {
+		Tool too = new McpSchema.Tool(TEST_TOOL_NAME, "Duplicate tool", emptyJsonSchema);
+
+		var mcpAsyncServer = McpServer.async(createMcpTransport())
+			.serverInfo("test-server", "1.0.0")
+			.capabilities(ServerCapabilities.builder().tools(true).build())
+			.tool(too, args -> Mono.just(new CallToolResult(List.of(), false)))
+			.build();
+
+		StepVerifier.create(mcpAsyncServer.notifyToolsListChanged()).verifyComplete();
+
+		assertThatCode(() -> mcpAsyncServer.closeGracefully().block(Duration.ofSeconds(10))).doesNotThrowAnyException();
+	}
+
+	// ---------------------------------------
+	// Resources Tests
+	// ---------------------------------------
+
+	@Test
+	void testNotifyResourcesListChanged() {
+		var mcpAsyncServer = McpServer.async(createMcpTransport()).serverInfo("test-server", "1.0.0").build();
+
+		StepVerifier.create(mcpAsyncServer.notifyResourcesListChanged()).verifyComplete();
+
+		assertThatCode(() -> mcpAsyncServer.closeGracefully().block(Duration.ofSeconds(10))).doesNotThrowAnyException();
+	}
+
+	@Test
+	void testAddResource() {
+		var mcpAsyncServer = McpServer.async(createMcpTransport())
+			.serverInfo("test-server", "1.0.0")
+			.capabilities(ServerCapabilities.builder().resources(true, false).build())
+			.build();
+
+		Resource resource = new Resource(TEST_RESOURCE_URI, "Test Resource", "text/plain", "Test resource description",
+				null);
+		McpServerFeatures.AsyncResourceRegistration registration = new McpServerFeatures.AsyncResourceRegistration(
+				resource, req -> Mono.just(new ReadResourceResult(List.of())));
+
+		StepVerifier.create(mcpAsyncServer.addResource(registration)).verifyComplete();
+
+		assertThatCode(() -> mcpAsyncServer.closeGracefully().block(Duration.ofSeconds(10))).doesNotThrowAnyException();
+	}
+
+	@Test
+	void testAddResourceWithNullRegistration() {
+		var mcpAsyncServer = McpServer.async(createMcpTransport())
+			.serverInfo("test-server", "1.0.0")
+			.capabilities(ServerCapabilities.builder().resources(true, false).build())
+			.build();
+
+		StepVerifier.create(mcpAsyncServer.addResource((McpServerFeatures.AsyncResourceRegistration) null))
+			.verifyErrorSatisfies(error -> {
+				assertThat(error).isInstanceOf(McpError.class).hasMessage("Resource must not be null");
+			});
+
+		assertThatCode(() -> mcpAsyncServer.closeGracefully().block(Duration.ofSeconds(10))).doesNotThrowAnyException();
+	}
+
+	@Test
+	void testAddResourceWithoutCapability() {
+		// Create a server without resource capabilities
+		McpAsyncServer serverWithoutResources = McpServer.async(createMcpTransport())
+			.serverInfo("test-server", "1.0.0")
+			.build();
+
+		Resource resource = new Resource(TEST_RESOURCE_URI, "Test Resource", "text/plain", "Test resource description",
+				null);
+		McpServerFeatures.AsyncResourceRegistration registration = new McpServerFeatures.AsyncResourceRegistration(
+				resource, req -> Mono.just(new ReadResourceResult(List.of())));
+
+		StepVerifier.create(serverWithoutResources.addResource(registration)).verifyErrorSatisfies(error -> {
+			assertThat(error).isInstanceOf(McpError.class)
+				.hasMessage("Server must be configured with resource capabilities");
+		});
+	}
+
+	@Test
+	void testRemoveResourceWithoutCapability() {
+		// Create a server without resource capabilities
+		McpAsyncServer serverWithoutResources = McpServer.async(createMcpTransport())
+			.serverInfo("test-server", "1.0.0")
+			.build();
+
+		StepVerifier.create(serverWithoutResources.removeResource(TEST_RESOURCE_URI)).verifyErrorSatisfies(error -> {
+			assertThat(error).isInstanceOf(McpError.class)
+				.hasMessage("Server must be configured with resource capabilities");
+		});
+	}
+
+	// ---------------------------------------
+	// Prompts Tests
+	// ---------------------------------------
+
+	@Test
+	void testNotifyPromptsListChanged() {
+		var mcpAsyncServer = McpServer.async(createMcpTransport()).serverInfo("test-server", "1.0.0").build();
+
+		StepVerifier.create(mcpAsyncServer.notifyPromptsListChanged()).verifyComplete();
+
+		assertThatCode(() -> mcpAsyncServer.closeGracefully().block(Duration.ofSeconds(10))).doesNotThrowAnyException();
+	}
+
+	@Test
+	void testAddPromptWithNullRegistration() {
+		var mcpAsyncServer = McpServer.async(createMcpTransport())
+			.serverInfo("test-server", "1.0.0")
+			.capabilities(ServerCapabilities.builder().prompts(false).build())
+			.build();
+
+		StepVerifier.create(mcpAsyncServer.addPrompt((McpServerFeatures.AsyncPromptRegistration) null))
+			.verifyErrorSatisfies(error -> {
+				assertThat(error).isInstanceOf(McpError.class).hasMessage("Prompt registration must not be null");
+			});
+	}
+
+	@Test
+	void testAddPromptWithoutCapability() {
+		// Create a server without prompt capabilities
+		McpAsyncServer serverWithoutPrompts = McpServer.async(createMcpTransport())
+			.serverInfo("test-server", "1.0.0")
+			.build();
+
+		Prompt prompt = new Prompt(TEST_PROMPT_NAME, "Test Prompt", List.of());
+		McpServerFeatures.AsyncPromptRegistration registration = new McpServerFeatures.AsyncPromptRegistration(prompt,
+				req -> Mono.just(new GetPromptResult("Test prompt description", List
+					.of(new PromptMessage(McpSchema.Role.ASSISTANT, new McpSchema.TextContent("Test content"))))));
+
+		StepVerifier.create(serverWithoutPrompts.addPrompt(registration)).verifyErrorSatisfies(error -> {
+			assertThat(error).isInstanceOf(McpError.class)
+				.hasMessage("Server must be configured with prompt capabilities");
+		});
+	}
+
+	@Test
+	void testRemovePromptWithoutCapability() {
+		// Create a server without prompt capabilities
+		McpAsyncServer serverWithoutPrompts = McpServer.async(createMcpTransport())
+			.serverInfo("test-server", "1.0.0")
+			.build();
+
+		StepVerifier.create(serverWithoutPrompts.removePrompt(TEST_PROMPT_NAME)).verifyErrorSatisfies(error -> {
+			assertThat(error).isInstanceOf(McpError.class)
+				.hasMessage("Server must be configured with prompt capabilities");
+		});
+	}
+
+	@Test
+	void testRemovePrompt() {
+		String TEST_PROMPT_NAME_TO_REMOVE = "TEST_PROMPT_NAME678";
+
+		Prompt prompt = new Prompt(TEST_PROMPT_NAME_TO_REMOVE, "Test Prompt", List.of());
+		McpServerFeatures.AsyncPromptRegistration registration = new McpServerFeatures.AsyncPromptRegistration(prompt,
+				req -> Mono.just(new GetPromptResult("Test prompt description", List
+					.of(new PromptMessage(McpSchema.Role.ASSISTANT, new McpSchema.TextContent("Test content"))))));
+
+		var mcpAsyncServer = McpServer.async(createMcpTransport())
+			.serverInfo("test-server", "1.0.0")
+			.capabilities(ServerCapabilities.builder().prompts(true).build())
+			.prompts(registration)
+			.build();
+
+		StepVerifier.create(mcpAsyncServer.removePrompt(TEST_PROMPT_NAME_TO_REMOVE)).verifyComplete();
+
+		assertThatCode(() -> mcpAsyncServer.closeGracefully().block(Duration.ofSeconds(10))).doesNotThrowAnyException();
+	}
+
+	@Test
+	void testRemoveNonexistentPrompt() {
+		var mcpAsyncServer2 = McpServer.async(createMcpTransport())
+			.serverInfo("test-server", "1.0.0")
+			.capabilities(ServerCapabilities.builder().prompts(true).build())
+			.build();
+
+		StepVerifier.create(mcpAsyncServer2.removePrompt("nonexistent-prompt")).verifyErrorSatisfies(error -> {
+			assertThat(error).isInstanceOf(McpError.class)
+				.hasMessage("Prompt with name 'nonexistent-prompt' not found");
+		});
+
+		assertThatCode(() -> mcpAsyncServer2.closeGracefully().block(Duration.ofSeconds(10)))
+			.doesNotThrowAnyException();
+	}
+
+	// ---------------------------------------
+	// Roots Tests
+	// ---------------------------------------
+
+	@Test
+	void testRootsChangeConsumers() {
+		// Test with single consumer
+		var rootsReceived = new McpSchema.Root[1];
+		var consumerCalled = new boolean[1];
+
+		var singleConsumerServer = McpServer.async(createMcpTransport())
+			.serverInfo("test-server", "1.0.0")
+			.rootsChangeConsumers(List.of(roots -> Mono.fromRunnable(() -> {
+				consumerCalled[0] = true;
+				if (!roots.isEmpty()) {
+					rootsReceived[0] = roots.get(0);
+				}
+			})))
+			.build();
+
+		assertThat(singleConsumerServer).isNotNull();
+		assertThatCode(() -> singleConsumerServer.closeGracefully().block(Duration.ofSeconds(10)))
+			.doesNotThrowAnyException();
+		onClose();
+
+		// Test with multiple consumers
+		var consumer1Called = new boolean[1];
+		var consumer2Called = new boolean[1];
+		var rootsContent = new List[1];
+
+		var multipleConsumersServer = McpServer.async(createMcpTransport())
+			.serverInfo("test-server", "1.0.0")
+			.rootsChangeConsumers(List.of(roots -> Mono.fromRunnable(() -> {
+				consumer1Called[0] = true;
+				rootsContent[0] = roots;
+			}), roots -> Mono.fromRunnable(() -> consumer2Called[0] = true)))
+			.build();
+
+		assertThat(multipleConsumersServer).isNotNull();
+		assertThatCode(() -> multipleConsumersServer.closeGracefully().block(Duration.ofSeconds(10)))
+			.doesNotThrowAnyException();
+		onClose();
+
+		// Test error handling
+		var errorHandlingServer = McpServer.async(createMcpTransport())
+			.serverInfo("test-server", "1.0.0")
+			.rootsChangeConsumers(List.of(roots -> {
+				throw new RuntimeException("Test error");
+			}))
+			.build();
+
+		assertThat(errorHandlingServer).isNotNull();
+		assertThatCode(() -> errorHandlingServer.closeGracefully().block(Duration.ofSeconds(10)))
+			.doesNotThrowAnyException();
+		onClose();
+
+		// Test without consumers
+		var noConsumersServer = McpServer.async(createMcpTransport()).serverInfo("test-server", "1.0.0").build();
+
+		assertThat(noConsumersServer).isNotNull();
+		assertThatCode(() -> noConsumersServer.closeGracefully().block(Duration.ofSeconds(10)))
+			.doesNotThrowAnyException();
+	}
+
+	// ---------------------------------------
+	// Logging Tests
+	// ---------------------------------------
+
+	@Test
+	void testLoggingLevels() {
+		var mcpAsyncServer = McpServer.async(createMcpTransport())
+			.serverInfo("test-server", "1.0.0")
+			.capabilities(ServerCapabilities.builder().logging().build())
+			.build();
+
+		// Test all logging levels
+		for (McpSchema.LoggingLevel level : McpSchema.LoggingLevel.values()) {
+			var notification = McpSchema.LoggingMessageNotification.builder()
+				.level(level)
+				.logger("test-logger")
+				.data("Test message with level " + level)
+				.build();
+
+			StepVerifier.create(mcpAsyncServer.loggingNotification(notification)).verifyComplete();
+		}
+	}
+
+	@Test
+	void testLoggingWithoutCapability() {
+		var mcpAsyncServer = McpServer.async(createMcpTransport())
+			.serverInfo("test-server", "1.0.0")
+			.capabilities(ServerCapabilities.builder().build()) // No logging capability
+			.build();
+
+		var notification = McpSchema.LoggingMessageNotification.builder()
+			.level(McpSchema.LoggingLevel.INFO)
+			.logger("test-logger")
+			.data("Test log message")
+			.build();
+
+		StepVerifier.create(mcpAsyncServer.loggingNotification(notification)).verifyComplete();
+	}
+
+	@Test
+	void testLoggingWithNullNotification() {
+		var mcpAsyncServer = McpServer.async(createMcpTransport())
+			.serverInfo("test-server", "1.0.0")
+			.capabilities(ServerCapabilities.builder().logging().build())
+			.build();
+
+		StepVerifier.create(mcpAsyncServer.loggingNotification(null)).verifyError(McpError.class);
+	}
+
+}

mcp-test/src/main/java/org/springframework/ai/mcp/server/AbstractMcpSyncServerTests.java

@@ -0,0 +1,444 @@
+/*
+ * 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 org.springframework.ai.mcp.server;
+
+import java.util.List;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import org.springframework.ai.mcp.spec.McpError;
+import org.springframework.ai.mcp.spec.McpSchema;
+import org.springframework.ai.mcp.spec.McpSchema.CallToolResult;
+import org.springframework.ai.mcp.spec.McpSchema.GetPromptResult;
+import org.springframework.ai.mcp.spec.McpSchema.Prompt;
+import org.springframework.ai.mcp.spec.McpSchema.PromptMessage;
+import org.springframework.ai.mcp.spec.McpSchema.ReadResourceResult;
+import org.springframework.ai.mcp.spec.McpSchema.Resource;
+import org.springframework.ai.mcp.spec.McpSchema.ServerCapabilities;
+import org.springframework.ai.mcp.spec.McpSchema.Tool;
+import org.springframework.ai.mcp.spec.McpTransport;
+import org.springframework.ai.mcp.spec.ServerMcpTransport;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatCode;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/**
+ * Test suite for the {@link McpSyncServer} that can be used with different
+ * {@link McpTransport} implementations.
+ *
+ * @author Christian Tzolov
+ */
+public abstract class AbstractMcpSyncServerTests {
+
+	private static final String TEST_TOOL_NAME = "test-tool";
+
+	private static final String TEST_RESOURCE_URI = "test://resource";
+
+	private static final String TEST_PROMPT_NAME = "test-prompt";
+
+	abstract protected ServerMcpTransport createMcpTransport();
+
+	protected void onStart() {
+	}
+
+	protected void onClose() {
+	}
+
+	@BeforeEach
+	void setUp() {
+		// onStart();
+	}
+
+	@AfterEach
+	void tearDown() {
+		onClose();
+	}
+
+	// ---------------------------------------
+	// Server Lifecycle Tests
+	// ---------------------------------------
+
+	@Test
+	void testConstructorWithInvalidArguments() {
+		assertThatThrownBy(() -> McpServer.sync(null)).isInstanceOf(IllegalArgumentException.class)
+			.hasMessage("Transport must not be null");
+
+		assertThatThrownBy(() -> McpServer.sync(createMcpTransport()).serverInfo(null))
+			.isInstanceOf(IllegalArgumentException.class)
+			.hasMessage("Server info must not be null");
+	}
+
+	@Test
+	void testGracefulShutdown() {
+		var mcpSyncServer = McpServer.sync(createMcpTransport()).serverInfo("test-server", "1.0.0").build();
+
+		assertThatCode(() -> mcpSyncServer.closeGracefully()).doesNotThrowAnyException();
+	}
+
+	@Test
+	void testImmediateClose() {
+		var mcpSyncServer = McpServer.sync(createMcpTransport()).serverInfo("test-server", "1.0.0").build();
+
+		assertThatCode(() -> mcpSyncServer.close()).doesNotThrowAnyException();
+	}
+
+	@Test
+	void testGetAsyncServer() {
+		var mcpSyncServer = McpServer.sync(createMcpTransport()).serverInfo("test-server", "1.0.0").build();
+
+		assertThat(mcpSyncServer.getAsyncServer()).isNotNull();
+
+		assertThatCode(() -> mcpSyncServer.closeGracefully()).doesNotThrowAnyException();
+	}
+
+	// ---------------------------------------
+	// Tools Tests
+	// ---------------------------------------
+
+	String emptyJsonSchema = """
+			{
+				"$schema": "http://json-schema.org/draft-07/schema#",
+				"type": "object",
+				"properties": {}
+			}
+			""";
+
+	@Test
+	void testAddTool() {
+		var mcpSyncServer = McpServer.sync(createMcpTransport())
+			.serverInfo("test-server", "1.0.0")
+			.capabilities(ServerCapabilities.builder().tools(true).build())
+			.build();
+
+		Tool newTool = new McpSchema.Tool("new-tool", "New test tool", emptyJsonSchema);
+		assertThatCode(() -> mcpSyncServer
+			.addTool(new McpServerFeatures.SyncToolRegistration(newTool, args -> new CallToolResult(List.of(), false))))
+			.doesNotThrowAnyException();
+
+		assertThatCode(() -> mcpSyncServer.closeGracefully()).doesNotThrowAnyException();
+	}
+
+	@Test
+	void testAddDuplicateTool() {
+		Tool duplicateTool = new McpSchema.Tool(TEST_TOOL_NAME, "Duplicate tool", emptyJsonSchema);
+
+		var mcpSyncServer = McpServer.sync(createMcpTransport())
+			.serverInfo("test-server", "1.0.0")
+			.capabilities(ServerCapabilities.builder().tools(true).build())
+			.tool(duplicateTool, args -> new CallToolResult(List.of(), false))
+			.build();
+
+		assertThatThrownBy(() -> mcpSyncServer.addTool(new McpServerFeatures.SyncToolRegistration(duplicateTool,
+				args -> new CallToolResult(List.of(), false))))
+			.isInstanceOf(McpError.class)
+			.hasMessage("Tool with name '" + TEST_TOOL_NAME + "' already exists");
+
+		assertThatCode(() -> mcpSyncServer.closeGracefully()).doesNotThrowAnyException();
+	}
+
+	@Test
+	void testRemoveTool() {
+		Tool tool = new McpSchema.Tool(TEST_TOOL_NAME, "Test tool", emptyJsonSchema);
+
+		var mcpSyncServer = McpServer.sync(createMcpTransport())
+			.serverInfo("test-server", "1.0.0")
+			.capabilities(ServerCapabilities.builder().tools(true).build())
+			.tool(tool, args -> new CallToolResult(List.of(), false))
+			.build();
+
+		assertThatCode(() -> mcpSyncServer.removeTool(TEST_TOOL_NAME)).doesNotThrowAnyException();
+
+		assertThatCode(() -> mcpSyncServer.closeGracefully()).doesNotThrowAnyException();
+	}
+
+	@Test
+	void testRemoveNonexistentTool() {
+		var mcpSyncServer = McpServer.sync(createMcpTransport())
+			.serverInfo("test-server", "1.0.0")
+			.capabilities(ServerCapabilities.builder().tools(true).build())
+			.build();
+
+		assertThatThrownBy(() -> mcpSyncServer.removeTool("nonexistent-tool")).isInstanceOf(McpError.class)
+			.hasMessage("Tool with name 'nonexistent-tool' not found");
+
+		assertThatCode(() -> mcpSyncServer.closeGracefully()).doesNotThrowAnyException();
+	}
+
+	@Test
+	void testNotifyToolsListChanged() {
+		var mcpSyncServer = McpServer.sync(createMcpTransport()).serverInfo("test-server", "1.0.0").build();
+
+		assertThatCode(() -> mcpSyncServer.notifyToolsListChanged()).doesNotThrowAnyException();
+
+		assertThatCode(() -> mcpSyncServer.closeGracefully()).doesNotThrowAnyException();
+	}
+
+	// ---------------------------------------
+	// Resources Tests
+	// ---------------------------------------
+
+	@Test
+	void testNotifyResourcesListChanged() {
+		var mcpSyncServer = McpServer.sync(createMcpTransport()).serverInfo("test-server", "1.0.0").build();
+
+		assertThatCode(() -> mcpSyncServer.notifyResourcesListChanged()).doesNotThrowAnyException();
+
+		assertThatCode(() -> mcpSyncServer.closeGracefully()).doesNotThrowAnyException();
+	}
+
+	@Test
+	void testAddResource() {
+		var mcpSyncServer = McpServer.sync(createMcpTransport())
+			.serverInfo("test-server", "1.0.0")
+			.capabilities(ServerCapabilities.builder().resources(true, false).build())
+			.build();
+
+		Resource resource = new Resource(TEST_RESOURCE_URI, "Test Resource", "text/plain", "Test resource description",
+				null);
+		McpServerFeatures.SyncResourceRegistration registration = new McpServerFeatures.SyncResourceRegistration(
+				resource, req -> new ReadResourceResult(List.of()));
+
+		assertThatCode(() -> mcpSyncServer.addResource(registration)).doesNotThrowAnyException();
+
+		assertThatCode(() -> mcpSyncServer.closeGracefully()).doesNotThrowAnyException();
+	}
+
+	@Test
+	void testAddResourceWithNullRegistration() {
+		var mcpSyncServer = McpServer.sync(createMcpTransport())
+			.serverInfo("test-server", "1.0.0")
+			.capabilities(ServerCapabilities.builder().resources(true, false).build())
+			.build();
+
+		assertThatThrownBy(() -> mcpSyncServer.addResource((McpServerFeatures.SyncResourceRegistration) null))
+			.isInstanceOf(McpError.class)
+			.hasMessage("Resource must not be null");
+
+		assertThatCode(() -> mcpSyncServer.closeGracefully()).doesNotThrowAnyException();
+	}
+
+	@Test
+	void testAddResourceWithoutCapability() {
+		var serverWithoutResources = McpServer.sync(createMcpTransport()).serverInfo("test-server", "1.0.0").build();
+
+		Resource resource = new Resource(TEST_RESOURCE_URI, "Test Resource", "text/plain", "Test resource description",
+				null);
+		McpServerFeatures.SyncResourceRegistration registration = new McpServerFeatures.SyncResourceRegistration(
+				resource, req -> new ReadResourceResult(List.of()));
+
+		assertThatThrownBy(() -> serverWithoutResources.addResource(registration)).isInstanceOf(McpError.class)
+			.hasMessage("Server must be configured with resource capabilities");
+	}
+
+	@Test
+	void testRemoveResourceWithoutCapability() {
+		var serverWithoutResources = McpServer.sync(createMcpTransport()).serverInfo("test-server", "1.0.0").build();
+
+		assertThatThrownBy(() -> serverWithoutResources.removeResource(TEST_RESOURCE_URI)).isInstanceOf(McpError.class)
+			.hasMessage("Server must be configured with resource capabilities");
+	}
+
+	// ---------------------------------------
+	// Prompts Tests
+	// ---------------------------------------
+
+	@Test
+	void testNotifyPromptsListChanged() {
+		var mcpSyncServer = McpServer.sync(createMcpTransport()).serverInfo("test-server", "1.0.0").build();
+
+		assertThatCode(() -> mcpSyncServer.notifyPromptsListChanged()).doesNotThrowAnyException();
+
+		assertThatCode(() -> mcpSyncServer.closeGracefully()).doesNotThrowAnyException();
+	}
+
+	@Test
+	void testAddPromptWithNullRegistration() {
+		var mcpSyncServer = McpServer.sync(createMcpTransport())
+			.serverInfo("test-server", "1.0.0")
+			.capabilities(ServerCapabilities.builder().prompts(false).build())
+			.build();
+
+		assertThatThrownBy(() -> mcpSyncServer.addPrompt((McpServerFeatures.SyncPromptRegistration) null))
+			.isInstanceOf(McpError.class)
+			.hasMessage("Prompt registration must not be null");
+	}
+
+	@Test
+	void testAddPromptWithoutCapability() {
+		var serverWithoutPrompts = McpServer.sync(createMcpTransport()).serverInfo("test-server", "1.0.0").build();
+
+		Prompt prompt = new Prompt(TEST_PROMPT_NAME, "Test Prompt", List.of());
+		McpServerFeatures.SyncPromptRegistration registration = new McpServerFeatures.SyncPromptRegistration(prompt,
+				req -> new GetPromptResult("Test prompt description", List
+					.of(new PromptMessage(McpSchema.Role.ASSISTANT, new McpSchema.TextContent("Test content")))));
+
+		assertThatThrownBy(() -> serverWithoutPrompts.addPrompt(registration)).isInstanceOf(McpError.class)
+			.hasMessage("Server must be configured with prompt capabilities");
+	}
+
+	@Test
+	void testRemovePromptWithoutCapability() {
+		var serverWithoutPrompts = McpServer.sync(createMcpTransport()).serverInfo("test-server", "1.0.0").build();
+
+		assertThatThrownBy(() -> serverWithoutPrompts.removePrompt(TEST_PROMPT_NAME)).isInstanceOf(McpError.class)
+			.hasMessage("Server must be configured with prompt capabilities");
+	}
+
+	@Test
+	void testRemovePrompt() {
+		Prompt prompt = new Prompt(TEST_PROMPT_NAME, "Test Prompt", List.of());
+		McpServerFeatures.SyncPromptRegistration registration = new McpServerFeatures.SyncPromptRegistration(prompt,
+				req -> new GetPromptResult("Test prompt description", List
+					.of(new PromptMessage(McpSchema.Role.ASSISTANT, new McpSchema.TextContent("Test content")))));
+
+		var mcpSyncServer = McpServer.sync(createMcpTransport())
+			.serverInfo("test-server", "1.0.0")
+			.capabilities(ServerCapabilities.builder().prompts(true).build())
+			.prompts(registration)
+			.build();
+
+		assertThatCode(() -> mcpSyncServer.removePrompt(TEST_PROMPT_NAME)).doesNotThrowAnyException();
+
+		assertThatCode(() -> mcpSyncServer.closeGracefully()).doesNotThrowAnyException();
+	}
+
+	@Test
+	void testRemoveNonexistentPrompt() {
+		var mcpSyncServer = McpServer.sync(createMcpTransport())
+			.serverInfo("test-server", "1.0.0")
+			.capabilities(ServerCapabilities.builder().prompts(true).build())
+			.build();
+
+		assertThatThrownBy(() -> mcpSyncServer.removePrompt("nonexistent-prompt")).isInstanceOf(McpError.class)
+			.hasMessage("Prompt with name 'nonexistent-prompt' not found");
+
+		assertThatCode(() -> mcpSyncServer.closeGracefully()).doesNotThrowAnyException();
+	}
+
+	// ---------------------------------------
+	// Roots Tests
+	// ---------------------------------------
+
+	@Test
+	void testRootsChangeConsumers() {
+		// Test with single consumer
+		var rootsReceived = new McpSchema.Root[1];
+		var consumerCalled = new boolean[1];
+
+		var singleConsumerServer = McpServer.sync(createMcpTransport())
+			.serverInfo("test-server", "1.0.0")
+			.rootsChangeConsumers(List.of(roots -> {
+				consumerCalled[0] = true;
+				if (!roots.isEmpty()) {
+					rootsReceived[0] = roots.get(0);
+				}
+			}))
+			.build();
+
+		assertThat(singleConsumerServer).isNotNull();
+		assertThatCode(() -> singleConsumerServer.closeGracefully()).doesNotThrowAnyException();
+		onClose();
+
+		// Test with multiple consumers
+		var consumer1Called = new boolean[1];
+		var consumer2Called = new boolean[1];
+		var rootsContent = new List[1];
+
+		var multipleConsumersServer = McpServer.sync(createMcpTransport())
+			.serverInfo("test-server", "1.0.0")
+			.rootsChangeConsumers(List.of(roots -> {
+				consumer1Called[0] = true;
+				rootsContent[0] = roots;
+			}, roots -> consumer2Called[0] = true))
+			.build();
+
+		assertThat(multipleConsumersServer).isNotNull();
+		assertThatCode(() -> multipleConsumersServer.closeGracefully()).doesNotThrowAnyException();
+		onClose();
+
+		// Test error handling
+		var errorHandlingServer = McpServer.sync(createMcpTransport())
+			.serverInfo("test-server", "1.0.0")
+			.rootsChangeConsumers(List.of(roots -> {
+				throw new RuntimeException("Test error");
+			}))
+			.build();
+
+		assertThat(errorHandlingServer).isNotNull();
+		assertThatCode(() -> errorHandlingServer.closeGracefully()).doesNotThrowAnyException();
+		onClose();
+
+		// Test without consumers
+		var noConsumersServer = McpServer.sync(createMcpTransport()).serverInfo("test-server", "1.0.0").build();
+
+		assertThat(noConsumersServer).isNotNull();
+		assertThatCode(() -> noConsumersServer.closeGracefully()).doesNotThrowAnyException();
+	}
+
+	// ---------------------------------------
+	// Logging Tests
+	// ---------------------------------------
+
+	@Test
+	void testLoggingLevels() {
+		var mcpSyncServer = McpServer.sync(createMcpTransport())
+			.serverInfo("test-server", "1.0.0")
+			.capabilities(ServerCapabilities.builder().logging().build())
+			.build();
+
+		// Test all logging levels
+		for (McpSchema.LoggingLevel level : McpSchema.LoggingLevel.values()) {
+			var notification = McpSchema.LoggingMessageNotification.builder()
+				.level(level)
+				.logger("test-logger")
+				.data("Test message with level " + level)
+				.build();
+
+			assertThatCode(() -> mcpSyncServer.loggingNotification(notification)).doesNotThrowAnyException();
+		}
+	}
+
+	@Test
+	void testLoggingWithoutCapability() {
+		var mcpSyncServer = McpServer.sync(createMcpTransport())
+			.serverInfo("test-server", "1.0.0")
+			.capabilities(ServerCapabilities.builder().build()) // No logging capability
+			.build();
+
+		var notification = McpSchema.LoggingMessageNotification.builder()
+			.level(McpSchema.LoggingLevel.INFO)
+			.logger("test-logger")
+			.data("Test log message")
+			.build();
+
+		assertThatCode(() -> mcpSyncServer.loggingNotification(notification)).doesNotThrowAnyException();
+	}
+
+	@Test
+	void testLoggingWithNullNotification() {
+		var mcpSyncServer = McpServer.sync(createMcpTransport())
+			.serverInfo("test-server", "1.0.0")
+			.capabilities(ServerCapabilities.builder().logging().build())
+			.build();
+
+		assertThatThrownBy(() -> mcpSyncServer.loggingNotification(null)).isInstanceOf(McpError.class);
+	}
+
+}

mcp-test/src/main/java/org/springframework/ai/mcp/server/transport/BlaBlockingInputStream.java

@@ -1,81 +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 org.springframework.ai.mcp.server.transport;
-
-import java.io.IOException;
-import java.io.InputStream;
-import java.util.concurrent.BlockingQueue;
-import java.util.concurrent.LinkedBlockingQueue;
-
-public class BlaBlockingInputStream extends InputStream {
-
-	private final BlockingQueue<Integer> queue = new LinkedBlockingQueue<>();
-
-	private volatile boolean completed = false;
-
-	private volatile boolean closed = false;
-
-	@Override
-	public int read() throws IOException {
-		if (closed) {
-			throw new IOException("Stream is closed");
-		}
-
-		try {
-			Integer value = queue.poll();
-			if (value == null) {
-				if (completed) {
-					return -1;
-				}
-				value = queue.take(); // Blocks until data is available
-				if (value == null && completed) {
-					return -1;
-				}
-			}
-			return value;
-		}
-		catch (InterruptedException e) {
-			Thread.currentThread().interrupt();
-			throw new IOException("Read interrupted", e);
-		}
-	}
-
-	public void write(int b) {
-		if (!closed && !completed) {
-			queue.offer(b);
-		}
-	}
-
-	public void write(byte[] data) {
-		if (!closed && !completed) {
-			for (byte b : data) {
-				queue.offer((int) b & 0xFF);
-			}
-		}
-	}
-
-	public void complete() {
-		this.completed = true;
-	}
-
-	@Override
-	public void close() {
-		this.closed = true;
-		this.completed = true;
-		this.queue.clear();
-	}
-
-}

mcp-transport/mcp-webflux-sse-transport/src/test/java/org/springframework/ai/mcp/client/WebFluxSseMcpAsyncClientTests.java

@@ -30,7 +30,7 @@ import org.springframework.web.reactive.function.client.WebClient;
  * @author Christian Tzolov
  */
 @Timeout(15) // Giving extra time beyond the client timeout
-class WebFluxSseMcpAsyncClientTests extends BlaAbstractMcpAsyncClientTests {
+class WebFluxSseMcpAsyncClientTests extends AbstractMcpAsyncClientTests {
 
 	static String host = "http://localhost:3001";
 

mcp-transport/mcp-webflux-sse-transport/src/test/java/org/springframework/ai/mcp/client/WebFluxSseMcpSyncClientTests.java

@@ -30,7 +30,7 @@ import org.springframework.web.reactive.function.client.WebClient;
  * @author Christian Tzolov
  */
 @Timeout(15) // Giving extra time beyond the client timeout
-class WebFluxSseMcpSyncClientTests extends BlaAbstractMcpSyncClientTests {
+class WebFluxSseMcpSyncClientTests extends AbstractMcpSyncClientTests {
 
 	static String host = "http://localhost:3001";
 

mcp-transport/mcp-webflux-sse-transport/src/test/java/org/springframework/ai/mcp/server/WebFluxSseMcpAsyncServerTests.java

@@ -33,7 +33,7 @@ import org.springframework.web.reactive.function.server.RouterFunctions;
  * @author Christian Tzolov
  */
 @Timeout(15) // Giving extra time beyond the client timeout
-class WebFluxSseMcpAsyncServerTests extends BlaAbstractMcpAsyncServerTests {
+class WebFluxSseMcpAsyncServerTests extends AbstractMcpAsyncServerTests {
 
 	private static final int PORT = 8181;
 

mcp-transport/mcp-webflux-sse-transport/src/test/java/org/springframework/ai/mcp/server/WebFluxSseMcpSyncServerTests.java

@@ -33,7 +33,7 @@ import org.springframework.web.reactive.function.server.RouterFunctions;
  * @author Christian Tzolov
  */
 @Timeout(15) // Giving extra time beyond the client timeout
-class WebFluxSseMcpSyncServerTests extends BlaAbstractMcpSyncServerTests {
+class WebFluxSseMcpSyncServerTests extends AbstractMcpSyncServerTests {
 
 	private static final int PORT = 8182;
 

mcp-transport/mcp-webmvc-sse-transport/src/test/java/org/springframework/ai/mcp/server/WebMvcSseAsyncServerTransportTests.java

@@ -33,7 +33,7 @@ import org.springframework.web.servlet.function.RouterFunction;
 import org.springframework.web.servlet.function.ServerResponse;
 
 @Timeout(15)
-class WebMvcSseAsyncServerTransportTests extends BlaAbstractMcpAsyncServerTests {
+class WebMvcSseAsyncServerTransportTests extends AbstractMcpAsyncServerTests {
 
 	private static final String MESSAGE_ENDPOINT = "/mcp/message";
 

mcp-transport/mcp-webmvc-sse-transport/src/test/java/org/springframework/ai/mcp/server/WebMvcSseSyncServerTransportTests.java

@@ -33,7 +33,7 @@ import org.springframework.web.servlet.function.RouterFunction;
 import org.springframework.web.servlet.function.ServerResponse;
 
 @Timeout(15)
-class WebMvcSseSyncServerTransportTests extends BlaAbstractMcpSyncServerTests {
+class WebMvcSseSyncServerTransportTests extends AbstractMcpSyncServerTests {
 
 	private static final String MESSAGE_ENDPOINT = "/mcp/message";
 

mcp/src/test/java/org/springframework/ai/mcp/client/AbstractMcpAsyncClientTests.java

@@ -54,6 +54,7 @@ import static org.assertj.core.api.Assertions.assertThatThrownBy;
  * @author Christian Tzolov
  * @author Dariusz Jędrzejczyk
  */
+// KEEP IN SYNC with the class in mcp-test module
 public abstract class AbstractMcpAsyncClientTests {
 
 	private McpAsyncClient mcpAsyncClient;

mcp/src/test/java/org/springframework/ai/mcp/client/AbstractMcpSyncClientTests.java

@@ -50,6 +50,7 @@ import static org.assertj.core.api.Assertions.assertThatThrownBy;
  * @author Christian Tzolov
  * @author Dariusz Jędrzejczyk
  */
+// KEEP IN SYNC with the class in mcp-test module
 public abstract class AbstractMcpSyncClientTests {
 
 	private McpSyncClient mcpSyncClient;

mcp/src/test/java/org/springframework/ai/mcp/server/AbstractMcpAsyncServerTests.java

@@ -48,6 +48,7 @@ import static org.assertj.core.api.Assertions.assertThatThrownBy;
  *
  * @author Christian Tzolov
  */
+// KEEP IN SYNC with the class in mcp-test module
 public abstract class AbstractMcpAsyncServerTests {
 
 	private static final String TEST_TOOL_NAME = "test-tool";

mcp/src/test/java/org/springframework/ai/mcp/server/AbstractMcpSyncServerTests.java

@@ -45,6 +45,7 @@ import static org.assertj.core.api.Assertions.assertThatThrownBy;
  *
  * @author Christian Tzolov
  */
+// KEEP IN SYNC with the class in mcp-test module
 public abstract class AbstractMcpSyncServerTests {
 
 	private static final String TEST_TOOL_NAME = "test-tool";