Refactor existing tests to have a shared base class

stephentoub · stephentoub · commit 8701faa92a9d · 2025-04-08T16:01:28.000-04:00
diff --git a/samples/EverythingServer/LoggingUpdateMessageSender.cs b/samples/EverythingServer/LoggingUpdateMessageSender.cs
@@ -1,5 +1,4 @@
-﻿using Microsoft.Extensions.DependencyInjection;
-using Microsoft.Extensions.Hosting;
+﻿using Microsoft.Extensions.Hosting;
 using ModelContextProtocol;
 using ModelContextProtocol.Protocol.Types;
 using ModelContextProtocol.Server;
diff --git a/samples/EverythingServer/ResourceGenerator.cs b/samples/EverythingServer/ResourceGenerator.cs
@@ -1,7 +1,4 @@
 using ModelContextProtocol.Protocol.Types;
-using System;
-using System.Collections.Generic;
-using System.Linq;
 
 namespace EverythingServer;
 
diff --git a/samples/EverythingServer/Tools/LongRunningTool.cs b/samples/EverythingServer/Tools/LongRunningTool.cs
@@ -1,5 +1,4 @@
 ﻿using ModelContextProtocol;
-using ModelContextProtocol.Protocol.Messages;
 using ModelContextProtocol.Protocol.Types;
 using ModelContextProtocol.Server;
 using System.ComponentModel;
diff --git a/src/ModelContextProtocol/Client/McpClientFactory.cs b/src/ModelContextProtocol/Client/McpClientFactory.cs
@@ -1,6 +1,4 @@
-﻿using System.Globalization;
-using System.Runtime.InteropServices;
-using ModelContextProtocol.Logging;
+﻿using ModelContextProtocol.Logging;
 using ModelContextProtocol.Protocol.Transport;
 using ModelContextProtocol.Utils;
 using Microsoft.Extensions.Logging;
diff --git a/src/ModelContextProtocol/Protocol/Types/ClientCapabilities.cs b/src/ModelContextProtocol/Protocol/Types/ClientCapabilities.cs
@@ -1,5 +1,4 @@
 ﻿using ModelContextProtocol.Protocol.Messages;
-using ModelContextProtocol.Server;
 using System.Text.Json.Serialization;
 
 namespace ModelContextProtocol.Protocol.Types;
diff --git a/src/ModelContextProtocol/Protocol/Types/LoggingCapability.cs b/src/ModelContextProtocol/Protocol/Types/LoggingCapability.cs
@@ -1,5 +1,4 @@
-﻿using ModelContextProtocol.Protocol.Messages;
-using ModelContextProtocol.Server;
+﻿using ModelContextProtocol.Server;
 using System.Text.Json.Serialization;
 
 namespace ModelContextProtocol.Protocol.Types;
diff --git a/src/ModelContextProtocol/Protocol/Types/PromptsCapability.cs b/src/ModelContextProtocol/Protocol/Types/PromptsCapability.cs
@@ -1,5 +1,4 @@
-﻿using ModelContextProtocol.Protocol.Messages;
-using ModelContextProtocol.Server;
+﻿using ModelContextProtocol.Server;
 using System.Text.Json.Serialization;
 
 namespace ModelContextProtocol.Protocol.Types;
diff --git a/src/ModelContextProtocol/Protocol/Types/ResourcesCapability.cs b/src/ModelContextProtocol/Protocol/Types/ResourcesCapability.cs
@@ -1,5 +1,4 @@
-﻿using ModelContextProtocol.Protocol.Messages;
-using ModelContextProtocol.Server;
+﻿using ModelContextProtocol.Server;
 using System.Text.Json.Serialization;
 
 namespace ModelContextProtocol.Protocol.Types;
diff --git a/src/ModelContextProtocol/Protocol/Types/RootsCapability.cs b/src/ModelContextProtocol/Protocol/Types/RootsCapability.cs
@@ -1,6 +1,4 @@
-﻿using ModelContextProtocol.Protocol.Messages;
-using ModelContextProtocol.Server;
-using System.Text.Json.Serialization;
+﻿using System.Text.Json.Serialization;
 
 namespace ModelContextProtocol.Protocol.Types;
 
diff --git a/src/ModelContextProtocol/Protocol/Types/SamplingCapability.cs b/src/ModelContextProtocol/Protocol/Types/SamplingCapability.cs
@@ -1,6 +1,4 @@
-﻿using ModelContextProtocol.Protocol.Messages;
-using ModelContextProtocol.Server;
-using System.Text.Json.Serialization;
+﻿using System.Text.Json.Serialization;
 
 namespace ModelContextProtocol.Protocol.Types;
 
diff --git a/src/ModelContextProtocol/Protocol/Types/ToolsCapability.cs b/src/ModelContextProtocol/Protocol/Types/ToolsCapability.cs
@@ -1,5 +1,4 @@
-﻿using ModelContextProtocol.Protocol.Messages;
-using ModelContextProtocol.Server;
+﻿using ModelContextProtocol.Server;
 using System.Text.Json.Serialization;
 
 namespace ModelContextProtocol.Protocol.Types;
diff --git a/src/ModelContextProtocol/Server/McpServerOptions.cs b/src/ModelContextProtocol/Server/McpServerOptions.cs
@@ -1,6 +1,5 @@
 ﻿
 using ModelContextProtocol.Protocol.Types;
-using System.Text.Json.Serialization;
 
 namespace ModelContextProtocol.Server;
 
diff --git a/src/ModelContextProtocol/Shared/McpSession.cs b/src/ModelContextProtocol/Shared/McpSession.cs
@@ -338,6 +338,8 @@ public async Task<JsonRpcResponse> SendRequestAsync(JsonRpcRequest request, Canc
             throw new McpException("Transport is not connected");
         }
 
+        cancellationToken.ThrowIfCancellationRequested();
+
         Histogram<double> durationMetric = _isServer ? s_serverRequestDuration : s_clientRequestDuration;
         string method = request.Method;
 
@@ -425,6 +427,8 @@ public async Task SendMessageAsync(IJsonRpcMessage message, CancellationToken ca
             throw new McpException("Transport is not connected");
         }
 
+        cancellationToken.ThrowIfCancellationRequested();
+
         Histogram<double> durationMetric = _isServer ? s_serverRequestDuration : s_clientRequestDuration;
         string method = GetMethodName(message);
 
diff --git a/tests/ModelContextProtocol.Tests/Client/McpClientExtensionsTests.cs b/tests/ModelContextProtocol.Tests/Client/McpClientExtensionsTests.cs
@@ -3,45 +3,32 @@
 using Microsoft.Extensions.Logging;
 using ModelContextProtocol.Client;
 using ModelContextProtocol.Protocol.Messages;
-using ModelContextProtocol.Protocol.Transport;
 using ModelContextProtocol.Protocol.Types;
 using ModelContextProtocol.Server;
-using ModelContextProtocol.Tests.Utils;
 using Moq;
-using System.IO.Pipelines;
 using System.Text.Json;
 using System.Text.Json.Serialization.Metadata;
 using System.Threading.Channels;
 
 namespace ModelContextProtocol.Tests.Client;
 
-public class McpClientExtensionsTests : LoggedTest
+public class McpClientExtensionsTests : ClientServerTestBase
 {
-    private readonly Pipe _clientToServerPipe = new();
-    private readonly Pipe _serverToClientPipe = new();
-    private readonly ServiceProvider _serviceProvider;
-    private readonly CancellationTokenSource _cts;
-    private readonly IMcpServer _server;
-    private readonly Task _serverTask;
-
     public McpClientExtensionsTests(ITestOutputHelper outputHelper)
         : base(outputHelper)
     {
-        ServiceCollection sc = new();
-        sc.AddSingleton(LoggerFactory);
-        sc.AddMcpServer().WithStreamServerTransport(_clientToServerPipe.Reader.AsStream(), _serverToClientPipe.Writer.AsStream());
+    }
+
+    protected override void ConfigureServices(ServiceCollection services, IMcpServerBuilder mcpServerBuilder)
+    {
         for (int f = 0; f < 10; f++)
         {
             string name = $"Method{f}";
-            sc.AddSingleton(McpServerTool.Create((int i) => $"{name} Result {i}", new() { Name = name }));
+            services.AddSingleton(McpServerTool.Create((int i) => $"{name} Result {i}", new() { Name = name }));
         }
-        sc.AddSingleton(McpServerTool.Create([McpServerTool(Destructive = false, OpenWorld = true)](string i) => $"{i} Result", new() { Name = "ValuesSetViaAttr" }));
-        sc.AddSingleton(McpServerTool.Create([McpServerTool(Destructive = false, OpenWorld = true)](string i) => $"{i} Result", new() { Name = "ValuesSetViaOptions", Destructive = true, OpenWorld = false, ReadOnly = true }));
-        _serviceProvider = sc.BuildServiceProvider();
+        services.AddSingleton(McpServerTool.Create([McpServerTool(Destructive = false, OpenWorld = true)] (string i) => $"{i} Result", new() { Name = "ValuesSetViaAttr" }));
+        services.AddSingleton(McpServerTool.Create([McpServerTool(Destructive = false, OpenWorld = true)] (string i) => $"{i} Result", new() { Name = "ValuesSetViaOptions", Destructive = true, OpenWorld = false, ReadOnly = true }));
 
-        _server = _serviceProvider.GetRequiredService<IMcpServer>();
-        _cts = CancellationTokenSource.CreateLinkedTokenSource(TestContext.Current.CancellationToken);
-        _serverTask = _server.RunAsync(cancellationToken: _cts.Token);
     }
 
     [Theory]
@@ -218,30 +205,6 @@ public async Task CreateSamplingHandler_ShouldHandleResourceMessages()
         Assert.Equal("endTurn", result.StopReason);
     }
 
-    public async ValueTask DisposeAsync()
-    {
-        await _cts.CancelAsync();
-
-        _clientToServerPipe.Writer.Complete();
-        _serverToClientPipe.Writer.Complete();
-
-        await _serverTask;
-
-        await _serviceProvider.DisposeAsync();
-        _cts.Dispose();
-    }
-
-    private async Task<IMcpClient> CreateMcpClientForServer()
-    {
-        return await McpClientFactory.CreateAsync(
-            new StreamClientTransport(
-                serverInput: _clientToServerPipe.Writer.AsStream(),
-                serverOutput: _serverToClientPipe.Reader.AsStream(),
-                LoggerFactory),
-            loggerFactory: LoggerFactory,
-            cancellationToken: TestContext.Current.CancellationToken);
-    }
-
     [Fact]
     public async Task ListToolsAsync_AllToolsReturned()
     {
@@ -377,15 +340,15 @@ public async Task AsClientLoggerProvider_MessagesSentToClient()
     {
         IMcpClient client = await CreateMcpClientForServer();
 
-        ILoggerProvider loggerProvider = _server.AsClientLoggerProvider();
+        ILoggerProvider loggerProvider = Server.AsClientLoggerProvider();
         Assert.Throws<ArgumentNullException>("categoryName", () => loggerProvider.CreateLogger(null!));
 
         ILogger logger = loggerProvider.CreateLogger("TestLogger");
         Assert.NotNull(logger);
 
         Assert.Null(logger.BeginScope(""));
 
-        Assert.Null(_server.LoggingLevel);
+        Assert.Null(Server.LoggingLevel);
         Assert.False(logger.IsEnabled(LogLevel.Trace));
         Assert.False(logger.IsEnabled(LogLevel.Debug));
         Assert.False(logger.IsEnabled(LogLevel.Information));
@@ -396,13 +359,13 @@ public async Task AsClientLoggerProvider_MessagesSentToClient()
         await client.SetLoggingLevel(LoggingLevel.Info, TestContext.Current.CancellationToken);
 
         DateTime start = DateTime.UtcNow;
-        while (_server.LoggingLevel is null)
+        while (Server.LoggingLevel is null)
         {
             await Task.Delay(1, TestContext.Current.CancellationToken);
             Assert.True(DateTime.UtcNow - start < TimeSpan.FromSeconds(10), "Timed out waiting for logging level to be set");
         }
 
-        Assert.Equal(LoggingLevel.Info, _server.LoggingLevel);
+        Assert.Equal(LoggingLevel.Info, Server.LoggingLevel);
         Assert.False(logger.IsEnabled(LogLevel.Trace));
         Assert.False(logger.IsEnabled(LogLevel.Debug));
         Assert.True(logger.IsEnabled(LogLevel.Information));
diff --git a/tests/ModelContextProtocol.Tests/ClientServerTestBase.cs b/tests/ModelContextProtocol.Tests/ClientServerTestBase.cs
@@ -0,0 +1,75 @@
+﻿using Microsoft.Extensions.AI;
+using Microsoft.Extensions.DependencyInjection;
+using ModelContextProtocol.Client;
+using ModelContextProtocol.Protocol.Transport;
+using ModelContextProtocol.Server;
+using ModelContextProtocol.Tests.Utils;
+using System.IO.Pipelines;
+
+namespace ModelContextProtocol.Tests;
+
+public abstract class ClientServerTestBase : LoggedTest, IAsyncDisposable
+{
+    private readonly Pipe _clientToServerPipe = new();
+    private readonly Pipe _serverToClientPipe = new();
+    private readonly IMcpServerBuilder _builder;
+    private readonly CancellationTokenSource _cts;
+    private readonly Task _serverTask;
+
+    public ClientServerTestBase(ITestOutputHelper testOutputHelper)
+        : base(testOutputHelper)
+    {
+        ServiceCollection sc = new();
+        sc.AddSingleton(LoggerFactory);
+        _builder = sc
+            .AddMcpServer()
+            .WithStreamServerTransport(_clientToServerPipe.Reader.AsStream(), _serverToClientPipe.Writer.AsStream());
+        ConfigureServices(sc, _builder);
+        ServiceProvider = sc.BuildServiceProvider();
+
+        _cts = CancellationTokenSource.CreateLinkedTokenSource(TestContext.Current.CancellationToken);
+        Server = ServiceProvider.GetRequiredService<IMcpServer>();
+        _serverTask = Server.RunAsync(_cts.Token);
+    }
+
+    protected IMcpServer Server { get; }
+
+    protected IServiceProvider ServiceProvider { get; }
+
+    protected virtual void ConfigureServices(ServiceCollection services, IMcpServerBuilder mcpServerBuilder)
+    {
+    }
+
+    public async ValueTask DisposeAsync()
+    {
+        await _cts.CancelAsync();
+
+        _clientToServerPipe.Writer.Complete();
+        _serverToClientPipe.Writer.Complete();
+
+        await _serverTask;
+
+        if (ServiceProvider is IAsyncDisposable asyncDisposable)
+        {
+            await asyncDisposable.DisposeAsync();
+        }
+        else if (ServiceProvider is IDisposable disposable)
+        {
+            disposable.Dispose();
+        }
+
+        _cts.Dispose();
+        Dispose();
+    }
+
+    protected async Task<IMcpClient> CreateMcpClientForServer(McpClientOptions? options = null)
+    {
+        return await McpClientFactory.CreateAsync(
+            new StreamClientTransport(
+                serverInput: _clientToServerPipe.Writer.AsStream(),
+                _serverToClientPipe.Reader.AsStream(),
+                LoggerFactory),
+            loggerFactory: LoggerFactory,
+            cancellationToken: TestContext.Current.CancellationToken);
+    }
+}
diff --git a/tests/ModelContextProtocol.Tests/Configuration/McpServerBuilderExtensionsPromptsTests.cs b/tests/ModelContextProtocol.Tests/Configuration/McpServerBuilderExtensionsPromptsTests.cs
diff --git a/tests/ModelContextProtocol.Tests/Configuration/McpServerBuilderExtensionsToolsTests.cs b/tests/ModelContextProtocol.Tests/Configuration/McpServerBuilderExtensionsToolsTests.cs
diff --git a/tests/ModelContextProtocol.Tests/Protocol/NotificationHandlerTests.cs b/tests/ModelContextProtocol.Tests/Protocol/NotificationHandlerTests.cs
diff --git a/tests/ModelContextProtocol.Tests/SseIntegrationTests.cs b/tests/ModelContextProtocol.Tests/SseIntegrationTests.cs

Original file line number	Diff line number	Diff line change
`@@ -1,5 +1,4 @@`
`1`	`1`	`using ModelContextProtocol;`
`2`		`-using ModelContextProtocol.Protocol.Messages;`
`3`	`2`	`using ModelContextProtocol.Protocol.Types;`
`4`	`3`	`using ModelContextProtocol.Server;`
`5`	`4`	`using System.ComponentModel;`
Original file line number	Diff line number	Diff line change
`@@ -1,5 +1,4 @@`
`1`	`1`	`using ModelContextProtocol.Protocol.Messages;`
`2`		`-using ModelContextProtocol.Server;`
`3`	`2`	`using System.Text.Json.Serialization;`
`4`	`3`
`5`	`4`	`namespace ModelContextProtocol.Protocol.Types;`