Skip to content

Commit f2d033d

Browse files
halter73halter73
committed
Try to reduce test flakiness
1 parent 94d5f18 commit f2d033d

8 files changed

Lines changed: 65 additions & 38 deletions

File tree

‎samples/AspNetCoreSseServer/McpEndpointRouteBuilderExtensions.cs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -26,9 +26,9 @@ public static IEndpointConventionBuilder MapMcpSse(this IEndpointRouteBuilder en
2626

2727
try
2828
{
29-
var serverTask = server.RunAsync(cancellationToken: requestAborted);
30-
await transport.RunAsync(cancellationToken: requestAborted);
31-
await serverTask;
29+
var transportTask = transport.RunAsync(cancellationToken: requestAborted);
30+
await server.RunAsync(cancellationToken: requestAborted);
31+
await transportTask;
3232
}
3333
catch (OperationCanceledException) when (requestAborted.IsCancellationRequested)
3434
{

‎src/ModelContextProtocol/Protocol/Transport/StdioClientStreamTransport.cs‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ internal sealed class StdioClientStreamTransport : TransportBase
2020
private readonly McpServerConfig _serverConfig;
2121
private readonly ILogger _logger;
2222
private readonly JsonSerializerOptions _jsonOptions;
23+
private readonly DataReceivedEventHandler _logProcessErrors;
2324
private Process? _process;
2425
private Task? _readTask;
2526
private CancellationTokenSource? _shutdownCts;
@@ -42,6 +43,7 @@ public StdioClientStreamTransport(StdioClientTransportOptions options, McpServer
4243
_options = options;
4344
_serverConfig = serverConfig;
4445
_logger = (ILogger?)loggerFactory?.CreateLogger<StdioClientTransport>() ?? NullLogger.Instance;
46+
_logProcessErrors = (sender, args) => _logger.TransportError(EndpointName, args.Data ?? "(no data)");
4547
_jsonOptions = McpJsonUtilities.DefaultOptions;
4648
}
4749

@@ -98,7 +100,7 @@ public async Task ConnectAsync(CancellationToken cancellationToken = default)
98100
_process = new Process { StartInfo = startInfo };
99101

100102
// Set up error logging
101-
_process.ErrorDataReceived += (sender, args) => _logger.TransportError(EndpointName, args.Data ?? "(no data)");
103+
_process.ErrorDataReceived += _logProcessErrors;
102104

103105
// We need both stdin and stdout to use a no-BOM UTF-8 encoding. On .NET Core,
104106
// we can use ProcessStartInfo.StandardOutputEncoding/StandardInputEncoding, but
@@ -277,6 +279,7 @@ private async Task CleanupAsync(CancellationToken cancellationToken)
277279
}
278280
finally
279281
{
282+
process.ErrorDataReceived -= _logProcessErrors;
280283
process.Dispose();
281284
_process = null;
282285
}

‎tests/ModelContextProtocol.Tests/ClientIntegrationTests.cs‎

Lines changed: 14 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ public async Task ConnectAndPing_Stdio(string clientId)
3939

4040
// Act
4141
await using var client = await _fixture.CreateClientAsync(clientId);
42-
await client.PingAsync(CancellationToken.None);
42+
await client.PingAsync(TestContext.Current.CancellationToken);
4343

4444
// Assert
4545
Assert.NotNull(client);
@@ -89,7 +89,7 @@ public async Task CallTool_Stdio_EchoServer(string clientId)
8989
{
9090
["message"] = "Hello MCP!"
9191
},
92-
CancellationToken.None
92+
TestContext.Current.CancellationToken
9393
);
9494

9595
// assert
@@ -141,7 +141,7 @@ public async Task GetPrompt_Stdio_SimplePrompt(string clientId)
141141

142142
// act
143143
await using var client = await _fixture.CreateClientAsync(clientId);
144-
var result = await client.GetPromptAsync("simple_prompt", null, CancellationToken.None);
144+
var result = await client.GetPromptAsync("simple_prompt", null, TestContext.Current.CancellationToken);
145145

146146
// assert
147147
Assert.NotNull(result);
@@ -161,7 +161,7 @@ public async Task GetPrompt_Stdio_ComplexPrompt(string clientId)
161161
{ "temperature", "0.7" },
162162
{ "style", "formal" }
163163
};
164-
var result = await client.GetPromptAsync("complex_prompt", arguments, CancellationToken.None);
164+
var result = await client.GetPromptAsync("complex_prompt", arguments, TestContext.Current.CancellationToken);
165165

166166
// assert
167167
Assert.NotNull(result);
@@ -177,7 +177,7 @@ public async Task GetPrompt_NonExistent_ThrowsException(string clientId)
177177
// act
178178
await using var client = await _fixture.CreateClientAsync(clientId);
179179
await Assert.ThrowsAsync<McpClientException>(() =>
180-
client.GetPromptAsync("non_existent_prompt", null, CancellationToken.None));
180+
client.GetPromptAsync("non_existent_prompt", null, TestContext.Current.CancellationToken));
181181
}
182182

183183
[Theory]
@@ -220,7 +220,7 @@ public async Task ReadResource_Stdio_TextResource(string clientId)
220220
await using var client = await _fixture.CreateClientAsync(clientId);
221221
// Odd numbered resources are text in the everything server (despite the docs saying otherwise)
222222
// 1 is index 0, which is "even" in the 0-based index
223-
var result = await client.ReadResourceAsync("test://static/resource/1", CancellationToken.None);
223+
var result = await client.ReadResourceAsync("test://static/resource/1", TestContext.Current.CancellationToken);
224224

225225
Assert.NotNull(result);
226226
Assert.Single(result.Contents);
@@ -237,7 +237,7 @@ public async Task ReadResource_Stdio_BinaryResource(string clientId)
237237
await using var client = await _fixture.CreateClientAsync(clientId);
238238
// Even numbered resources are binary in the everything server (despite the docs saying otherwise)
239239
// 2 is index 1, which is "odd" in the 0-based index
240-
var result = await client.ReadResourceAsync("test://static/resource/2", CancellationToken.None);
240+
var result = await client.ReadResourceAsync("test://static/resource/2", TestContext.Current.CancellationToken);
241241

242242
Assert.NotNull(result);
243243
Assert.Single(result.Contents);
@@ -260,7 +260,7 @@ public async Task SubscribeResource_Stdio()
260260
tcs.TrySetResult(true);
261261
return Task.CompletedTask;
262262
});
263-
await client.SubscribeToResourceAsync("test://static/resource/1", CancellationToken.None);
263+
await client.SubscribeToResourceAsync("test://static/resource/1", TestContext.Current.CancellationToken);
264264

265265
await tcs.Task;
266266
}
@@ -281,13 +281,13 @@ public async Task UnsubscribeResource_Stdio()
281281
receivedNotification.TrySetResult(true);
282282
return Task.CompletedTask;
283283
});
284-
await client.SubscribeToResourceAsync("test://static/resource/1", CancellationToken.None);
284+
await client.SubscribeToResourceAsync("test://static/resource/1", TestContext.Current.CancellationToken);
285285

286286
// wait until we received a notification
287287
await receivedNotification.Task;
288288

289289
// unsubscribe
290-
await client.UnsubscribeFromResourceAsync("test://static/resource/1", CancellationToken.None);
290+
await client.UnsubscribeFromResourceAsync("test://static/resource/1", TestContext.Current.CancellationToken);
291291
receivedNotification = new();
292292

293293
// wait a bit to validate we don't receive another. this is best effort only;
@@ -309,7 +309,7 @@ public async Task GetCompletion_Stdio_ResourceReference(string clientId)
309309
Uri = "test://static/resource/1"
310310
},
311311
"argument_name", "1",
312-
CancellationToken.None
312+
TestContext.Current.CancellationToken
313313
);
314314

315315
Assert.NotNull(result);
@@ -331,7 +331,7 @@ public async Task GetCompletion_Stdio_PromptReference(string clientId)
331331
Name = "irrelevant"
332332
},
333333
argumentName: "style", argumentValue: "fo",
334-
CancellationToken.None
334+
TestContext.Current.CancellationToken
335335
);
336336

337337
Assert.NotNull(result);
@@ -411,7 +411,7 @@ public async Task Sampling_Stdio(string clientId)
411411
// });
412412

413413
// // Connect
414-
// await client.ConnectAsync(CancellationToken.None);
414+
// await client.ConnectAsync(TestContext.Current.CancellationToken);
415415

416416
// // assert
417417
// // nothing to assert, no servers implement roots, so we if no exception is thrown, it's a success
@@ -560,7 +560,7 @@ public async Task SetLoggingLevel_ReceivesLoggingMessages(string clientId)
560560
});
561561

562562
// act
563-
await client.SetLoggingLevel(LoggingLevel.Debug, CancellationToken.None);
563+
await client.SetLoggingLevel(LoggingLevel.Debug, TestContext.Current.CancellationToken);
564564

565565
// assert
566566
await receivedNotification.Task;

‎tests/ModelContextProtocol.Tests/Server/McpServerTests.cs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -538,7 +538,7 @@ await transport.SendMessageAsync(
538538
}
539539
);
540540

541-
var response = await receivedMessage.Task.WaitAsync(TimeSpan.FromSeconds(1));
541+
var response = await receivedMessage.Task.WaitAsync(TimeSpan.FromSeconds(5));
542542
Assert.NotNull(response);
543543
Assert.NotNull(response.Result);
544544

‎tests/ModelContextProtocol.Tests/SseServerIntegrationTestFixture.cs‎

Lines changed: 7 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
using ModelContextProtocol.Configuration;
44
using ModelContextProtocol.Protocol.Transport;
55
using ModelContextProtocol.Test.Utils;
6+
using ModelContextProtocol.Tests.Utils;
67
using ModelContextProtocol.TestSseServer;
78

89
namespace ModelContextProtocol.Tests;
@@ -47,6 +48,11 @@ public void Initialize(ITestOutputHelper output)
4748
_delegatingTestOutputHelper.CurrentTestOutputHelper = output;
4849
}
4950

51+
public void TestCompleted()
52+
{
53+
_delegatingTestOutputHelper.CurrentTestOutputHelper = null;
54+
}
55+
5056
public async ValueTask DisposeAsync()
5157
{
5258
_delegatingTestOutputHelper.CurrentTestOutputHelper = null;
@@ -61,16 +67,4 @@ public async ValueTask DisposeAsync()
6167
_redirectingLoggerFactory.Dispose();
6268
_stopCts.Dispose();
6369
}
64-
65-
private class DelegatingTestOutputHelper() : ITestOutputHelper
66-
{
67-
public ITestOutputHelper? CurrentTestOutputHelper { get; set; }
68-
69-
public string Output => CurrentTestOutputHelper?.Output ?? string.Empty;
70-
71-
public void Write(string message) => CurrentTestOutputHelper?.Write(message);
72-
public void Write(string format, params object[] args) => CurrentTestOutputHelper?.Write(format, args);
73-
public void WriteLine(string message) => CurrentTestOutputHelper?.WriteLine(message);
74-
public void WriteLine(string format, params object[] args) => CurrentTestOutputHelper?.WriteLine(format, args);
75-
}
76-
}
70+
}

‎tests/ModelContextProtocol.Tests/SseServerIntegrationTests.cs‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,12 @@ public SseServerIntegrationTests(SseServerIntegrationTestFixture fixture, ITestO
1515
_fixture.Initialize(testOutputHelper);
1616
}
1717

18+
public override void Dispose()
19+
{
20+
_fixture.TestCompleted();
21+
base.Dispose();
22+
}
23+
1824
private Task<IMcpClient> GetClientAsync(McpClientOptions? options = null)
1925
{
2026
return McpClientFactory.CreateAsync(

‎tests/ModelContextProtocol.Tests/Utils/DelegatingTestOutputHelper.cs‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
namespace ModelContextProtocol.Tests.Utils;
2+
3+
public class DelegatingTestOutputHelper() : ITestOutputHelper
4+
{
5+
public ITestOutputHelper? CurrentTestOutputHelper { get; set; }
6+
7+
public string Output => CurrentTestOutputHelper?.Output ?? string.Empty;
8+
9+
public void Write(string message) => CurrentTestOutputHelper?.Write(message);
10+
public void Write(string format, params object[] args) => CurrentTestOutputHelper?.Write(format, args);
11+
public void WriteLine(string message) => CurrentTestOutputHelper?.WriteLine(message);
12+
public void WriteLine(string format, params object[] args) => CurrentTestOutputHelper?.WriteLine(format, args);
13+
}

‎tests/ModelContextProtocol.Tests/Utils/LoggedTest.cs‎

Lines changed: 17 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -3,16 +3,27 @@
33

44
namespace ModelContextProtocol.Tests.Utils;
55

6-
public class LoggedTest(ITestOutputHelper testOutputHelper)
6+
public class LoggedTest : IDisposable
77
{
8-
public ITestOutputHelper TestOutputHelper { get; } = testOutputHelper;
9-
public ILoggerFactory LoggerFactory { get; } = CreateLoggerFactory(testOutputHelper);
8+
private readonly DelegatingTestOutputHelper _delegatingTestOutputHelper;
109

11-
private static ILoggerFactory CreateLoggerFactory(ITestOutputHelper testOutputHelper)
10+
public LoggedTest(ITestOutputHelper testOutputHelper)
1211
{
13-
return Microsoft.Extensions.Logging.LoggerFactory.Create(builder =>
12+
_delegatingTestOutputHelper = new()
1413
{
15-
builder.AddProvider(new XunitLoggerProvider(testOutputHelper));
14+
CurrentTestOutputHelper = testOutputHelper,
15+
};
16+
LoggerFactory = Microsoft.Extensions.Logging.LoggerFactory.Create(builder =>
17+
{
18+
builder.AddProvider(new XunitLoggerProvider(_delegatingTestOutputHelper));
1619
});
1720
}
21+
22+
public ITestOutputHelper TestOutputHelper => _delegatingTestOutputHelper;
23+
public ILoggerFactory LoggerFactory { get; }
24+
25+
public virtual void Dispose()
26+
{
27+
_delegatingTestOutputHelper.CurrentTestOutputHelper = null;
28+
}
1829
}

0 commit comments

Comments
 (0)