Skip to content

Commit 90536c2

Browse files
Tyler-R-KendrickTyler-R-Kendrick
committed
added tests
1 parent 8dc1f5d commit 90536c2

File tree

6 files changed

+361
-21
lines changed

6 files changed

+361
-21
lines changed

‎src/ModelContextProtocol/Client/IMcpClient.cs

+2-21
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,13 @@
11
using ModelContextProtocol.Protocol.Messages;
22
using ModelContextProtocol.Protocol.Types;
3+
using ModelContextProtocol.Shared;
34

45
namespace ModelContextProtocol.Client;
56

67
/// <summary>
78
/// Represents an instance of an MCP client connecting to a specific server.
89
/// </summary>
9-
public interface IMcpClient : IAsyncDisposable
10+
public interface IMcpClient : IMcpSession, IAsyncDisposable
1011
{
1112
/// <summary>
1213
/// Gets the capabilities supported by the server.
@@ -40,24 +41,4 @@ public interface IMcpClient : IAsyncDisposable
4041
/// </para>
4142
/// </remarks>
4243
void AddNotificationHandler(string method, Func<JsonRpcNotification, Task> handler);
43-
44-
/// <summary>
45-
/// Sends a generic JSON-RPC request to the server.
46-
/// </summary>
47-
/// <typeparam name="TResult">The expected response type.</typeparam>
48-
/// <param name="request">The JSON-RPC request to send.</param>
49-
/// <param name="cancellationToken">A token to cancel the operation.</param>
50-
/// <returns>A task containing the server's response.</returns>
51-
/// <remarks>
52-
/// It is recommended to use the capability-specific methods that use this one in their implementation.
53-
/// Use this method for custom requests or those not yet covered explicitly.
54-
/// </remarks>
55-
Task<TResult> SendRequestAsync<TResult>(JsonRpcRequest request, CancellationToken cancellationToken = default) where TResult : class;
56-
57-
/// <summary>
58-
/// Sends a message to the server.
59-
/// </summary>
60-
/// <param name="message">The message.</param>
61-
/// <param name="cancellationToken">A token to cancel the operation.</param>
62-
Task SendMessageAsync(IJsonRpcMessage message, CancellationToken cancellationToken = default);
6344
}

‎src/ModelContextProtocol/Shared/IMcpSession.cs

+26
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
using ModelContextProtocol.Protocol.Messages;
2+
3+
namespace ModelContextProtocol.Shared;
4+
5+
/// <summary>
6+
/// Class for managing an MCP JSON-RPC session. This covers both MCP clients and servers.
7+
/// </summary>
8+
public interface IMcpSession : IAsyncDisposable
9+
{
10+
/// <summary>
11+
/// Sends a generic JSON-RPC request to the server.
12+
/// </summary>
13+
/// <param name="message">The JSON-RPC request to send.</param>
14+
/// <param name="cancellationToken">A token to cancel the operation.</param>
15+
/// <returns>A task containing the server's response.</returns>
16+
Task SendMessageAsync(IJsonRpcMessage message, CancellationToken cancellationToken = default);
17+
18+
/// <summary>
19+
/// Sends a request over the protocol
20+
/// </summary>
21+
/// <typeparam name="TResult">The MCP Response type.</typeparam>
22+
/// <param name="request">The request instance</param>
23+
/// <param name="cancellationToken">The token for cancellation.</param>
24+
/// <returns>The MCP response.</returns>
25+
Task<TResult> SendRequestAsync<TResult>(JsonRpcRequest request, CancellationToken cancellationToken = default) where TResult : class;
26+
}

‎src/ModelContextProtocol/Shared/McpSessionExtensions.cs

+101
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,101 @@
1+
using ModelContextProtocol.Protocol.Messages;
2+
using ModelContextProtocol.Shared;
3+
using ModelContextProtocol.Utils;
4+
5+
namespace ModelContextProtocol;
6+
7+
/// <summary>Provides extension methods for interacting with an <see cref="IMcpSession"/>.</summary>
8+
public static class McpEndpointExtensions
9+
{
10+
/// <summary>
11+
/// Notifies the connected endpoint of an event.
12+
/// </summary>
13+
/// <param name="endpoint">The endpoint issuing the notification.</param>
14+
/// <param name="notification">The notification to send.</param>
15+
/// <param name="cancellationToken">A token to cancel the operation.</param>
16+
/// <exception cref="ArgumentNullException"><paramref name="endpoint"/> is <see langword="null"/>.</exception>
17+
/// <returns>A task representing the completion of the operation.</returns>
18+
public static Task NotifyAsync(
19+
this IMcpSession endpoint,
20+
JsonRpcNotification notification,
21+
CancellationToken cancellationToken = default)
22+
{
23+
Throw.IfNull(endpoint);
24+
25+
return endpoint.SendMessageAsync(notification, cancellationToken);
26+
}
27+
28+
/// <summary>
29+
/// Notifies the connected endpoint of an event.
30+
/// </summary>
31+
/// <param name="endpoint">The endpoint issuing the notification.</param>
32+
/// <param name="method">The method to call.</param>
33+
/// <param name="parameters">The parameters to send.</param>
34+
/// <param name="cancellationToken">A token to cancel the operation.</param>
35+
/// <exception cref="ArgumentNullException"><paramref name="endpoint"/> is <see langword="null"/>.</exception>
36+
/// <returns>A task representing the completion of the operation.</returns>
37+
public static Task NotifyAsync(
38+
this IMcpSession endpoint,
39+
string method,
40+
object? parameters = null,
41+
CancellationToken cancellationToken = default)
42+
{
43+
Throw.IfNull(endpoint);
44+
45+
return endpoint.NotifyAsync(new()
46+
{
47+
Method = method,
48+
Params = parameters,
49+
}, cancellationToken);
50+
}
51+
52+
/// <summary>Notifies the connected endpoint of progress.</summary>
53+
/// <param name="endpoint">The endpoint issuing the notification.</param>
54+
/// <param name="progressToken">The <see cref="ProgressToken"/> identifying the operation.</param>
55+
/// <param name="progress">The progress update to send.</param>
56+
/// <param name="cancellationToken">A token to cancel the operation.</param>
57+
/// <returns>A task representing the completion of the operation.</returns>
58+
/// <exception cref="ArgumentNullException"><paramref name="endpoint"/> is <see langword="null"/>.</exception>
59+
public static Task NotifyProgressAsync(
60+
this IMcpSession endpoint,
61+
ProgressToken progressToken,
62+
ProgressNotificationValue progress,
63+
CancellationToken cancellationToken = default)
64+
{
65+
Throw.IfNull(endpoint);
66+
67+
return endpoint.NotifyAsync(
68+
NotificationMethods.ProgressNotification,
69+
new ProgressNotification()
70+
{
71+
ProgressToken = progressToken,
72+
Progress = progress,
73+
}, cancellationToken);
74+
}
75+
76+
/// <summary>
77+
/// Notifies the connected endpoint that a request has been cancelled.
78+
/// </summary>
79+
/// <param name="endpoint">The endpoint issuing the notification.</param>
80+
/// <param name="requestId">The ID of the request to cancel.</param>
81+
/// <param name="reason">An optional reason for the cancellation.</param>
82+
/// <param name="cancellationToken">A token to cancel the operation.</param>
83+
/// <returns>A task representing the completion of the operation.</returns>
84+
/// <exception cref="ArgumentNullException"><paramref name="endpoint"/> is <see langword="null"/>.</exception>
85+
public static Task NotifyCancelAsync(
86+
this IMcpSession endpoint,
87+
RequestId requestId,
88+
string? reason = null,
89+
CancellationToken cancellationToken = default)
90+
{
91+
Throw.IfNull(endpoint);
92+
93+
return endpoint.NotifyAsync(
94+
NotificationMethods.CancelledNotification,
95+
new CancelledNotification()
96+
{
97+
RequestId = requestId,
98+
Reason = reason,
99+
}, cancellationToken);
100+
}
101+
}

‎tests/ModelContextProtocol.Tests/CancelledNotificationTests.cs

+91
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,91 @@
1+
using ModelContextProtocol.Protocol.Messages;
2+
using ModelContextProtocol.Protocol.Types;
3+
using ModelContextProtocol.Tests.Utils;
4+
5+
namespace ModelContextProtocol.Tests;
6+
7+
/// <summary>
8+
/// Tests for the cancelled notifications against an IMcpEndpoint.
9+
/// </summary>
10+
public class CancelledNotificationTests(
11+
McpEndpointTestFixture fixture, ITestOutputHelper testOutputHelper)
12+
: LoggedTest(testOutputHelper), IClassFixture<McpEndpointTestFixture>
13+
{
14+
[Fact]
15+
public async Task NotifyCancelAsync_SendsCorrectNotification()
16+
{
17+
// Arrange
18+
var token = TestContext.Current.CancellationToken;
19+
var clientTransport = fixture.CreateClientTransport();
20+
await using var endpoint = await fixture.CreateClientEndpointAsync(clientTransport);
21+
var transport = await clientTransport.ConnectAsync(token);
22+
23+
var requestId = new RequestId("test-request-id-123");
24+
const string reason = "Operation was cancelled by the user";
25+
26+
// Act
27+
await endpoint.NotifyCancelAsync(requestId, reason, token);
28+
29+
// Assert
30+
Assert.Equal(1, transport.MessageReader.Count);
31+
var message = await transport.MessageReader.ReadAsync(token);
32+
Assert.NotNull(message);
33+
34+
var notification = Assert.IsType<JsonRpcNotification>(message);
35+
Assert.Equal(NotificationMethods.CancelledNotification, notification.Method);
36+
37+
var cancelParams = Assert.IsType<CancelledNotification>(notification.Params);
38+
Assert.Equal(requestId, cancelParams.RequestId);
39+
Assert.Equal(reason, cancelParams.Reason);
40+
}
41+
42+
[Fact]
43+
public async Task SendRequestAsync_Cancellation_SendsNotification()
44+
{
45+
// Arrange
46+
var token = TestContext.Current.CancellationToken;
47+
var clientTransport = fixture.CreateClientTransport();
48+
await using var endpoint = await fixture.CreateClientEndpointAsync(clientTransport);
49+
var transport = await clientTransport.ConnectAsync(token);
50+
var requestId = new RequestId("test-request-id-123");
51+
JsonRpcRequest request = new()
52+
{
53+
Id = requestId,
54+
Method = "test.method",
55+
Params = new { },
56+
};
57+
using CancellationTokenSource cancellationSource = new();
58+
await cancellationSource.CancelAsync();
59+
// Act
60+
try
61+
{
62+
await endpoint.SendRequestAsync<EmptyResult>(request, cancellationSource.Token);
63+
}
64+
catch (OperationCanceledException)
65+
{
66+
// Expected exception
67+
}
68+
catch (Exception ex)
69+
{
70+
Assert.Fail($"Unexpected exception: {ex.Message}");
71+
}
72+
73+
// Assert
74+
Assert.Equal(2, transport.MessageReader.Count);
75+
var message = await transport.MessageReader.ReadAsync(token);
76+
Assert.NotNull(message);
77+
78+
var notification = Assert.IsType<JsonRpcNotification>(message);
79+
Assert.Equal(NotificationMethods.CancelledNotification, notification.Method);
80+
81+
var cancelParams = Assert.IsType<CancelledNotification>(notification.Params);
82+
Assert.Equal(requestId, cancelParams.RequestId);
83+
84+
message = await transport.MessageReader.ReadAsync(token);
85+
Assert.NotNull(message);
86+
var requestMessage = Assert.IsType<JsonRpcRequest>(message);
87+
Assert.Equal(request.Id, requestMessage.Id);
88+
Assert.Equal(request.Method, requestMessage.Method);
89+
Assert.Equal(request.Params, requestMessage.Params);
90+
}
91+
}

‎tests/ModelContextProtocol.Tests/McpEndpointTestFixture.cs

+80
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,80 @@
1+
using ModelContextProtocol.Client;
2+
using ModelContextProtocol.Server;
3+
using ModelContextProtocol.Protocol.Messages;
4+
using ModelContextProtocol.Protocol.Transport;
5+
using Microsoft.Extensions.Logging;
6+
using System.Threading.Channels;
7+
8+
namespace ModelContextProtocol.Tests;
9+
10+
/// <summary>
11+
/// Test fixture for McpEndpoint tests that provides shared transport implementations.
12+
/// </summary>
13+
public class McpEndpointTestFixture() : IAsyncDisposable
14+
{
15+
/// <summary>
16+
/// Creates a test transport.
17+
/// </summary>
18+
internal TestCancellationTransport CreateTransport() => new();
19+
internal IClientTransport CreateClientTransport(ITransport? transport = default)
20+
=> new TestCancellationClientTransport(transport ?? CreateTransport());
21+
22+
23+
/// <summary>
24+
/// Creates a test client endpoint.
25+
/// </summary>
26+
internal async Task<IMcpClient> CreateClientEndpointAsync(
27+
Func<McpServerConfig, ILoggerFactory?, IClientTransport>? transportFactory = default)
28+
{
29+
transportFactory ??= (_, _) => CreateClientTransport();
30+
return await McpClientFactory.CreateAsync(new()
31+
{
32+
Id = "TestServer",
33+
Name = "Test Server",
34+
TransportType = "TestTransport",
35+
}, createTransportFunc: transportFactory);
36+
}
37+
internal Task<IMcpClient> CreateClientEndpointAsync(IClientTransport transport)
38+
=> CreateClientEndpointAsync((_, _) => transport);
39+
internal Task<IMcpClient> CreateClientEndpointAsync(ITransport transport)
40+
=> CreateClientEndpointAsync(new TestCancellationClientTransport(transport));
41+
42+
internal async Task<IMcpServer> CreateServerEndpointAsync(ITransport transport)
43+
{
44+
var server = McpServerFactory.Create(transport, new()
45+
{
46+
ServerInfo = new()
47+
{
48+
Name = "TestServer",
49+
Version = "1.0.0",
50+
}
51+
});
52+
await server.RunAsync();
53+
return server;
54+
}
55+
56+
public ValueTask DisposeAsync() => ValueTask.CompletedTask;
57+
58+
internal class TestCancellationTransport : ITransport
59+
{
60+
public bool IsConnected => true;
61+
public List<IJsonRpcMessage> SentMessages { get; } = [];
62+
public ChannelReader<IJsonRpcMessage> MessageReader { get; init; }
63+
= Channel.CreateUnbounded<IJsonRpcMessage>().Reader;
64+
public Task SendMessageAsync(IJsonRpcMessage message, CancellationToken cancellationToken)
65+
{
66+
SentMessages.Add(message);
67+
return Task.CompletedTask;
68+
}
69+
70+
public ValueTask DisposeAsync() => ValueTask.CompletedTask;
71+
}
72+
73+
internal class TestCancellationClientTransport(ITransport transport) : IClientTransport
74+
{
75+
public Task<ITransport> ConnectAsync(CancellationToken cancellationToken = default)
76+
=> Task.FromResult(transport);
77+
78+
public ValueTask DisposeAsync() => ValueTask.CompletedTask;
79+
}
80+
}

0 commit comments

Comments
 (0)