Reduce absolute time dependencies in tests (#1649)

This commit is contained in:
Pranav Senthilnathan 2026-06-24 10:20:09 -07:00 • committed by GitHub
commit e7b450efc1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 123 additions and 80 deletions

View file

@ -1,7 +1,6 @@
using ModelContextProtocol.Protocol; using ModelContextProtocol.Protocol;
using ModelContextProtocol.Server; using ModelContextProtocol.Server;
using Microsoft.Extensions.DependencyInjection; using Microsoft.Extensions.DependencyInjection;
using System.Collections.Concurrent;
using System.Runtime.InteropServices; using System.Runtime.InteropServices;
using System.Text.Json; using System.Text.Json;
using System.Text.Json.Nodes; using System.Text.Json.Nodes;
@ -179,17 +178,25 @@ public class McpServerTaskTests : ClientServerTestBase
await using var client = await CreateMcpClientForServer(); await using var client = await CreateMcpClientForServer();
var ct = TestContext.Current.CancellationToken; var ct = TestContext.Current.CancellationToken;
_ = Task.Run(async () => var failedTask = new TaskCompletionSource<bool>();
// Run failure task once the task from the tool call is created
_taskStore.OnTaskCreated += taskId =>
{ {
await Task.Delay(100, ct); _ = Task.Run(async () =>
var taskId = _taskStore.GetAllTaskIds().Single(); {
_taskStore.FailTask(taskId, JsonElement.Parse("""{"code":-32000,"message":"something went wrong"}""")); await Task.Delay(100, ct);
}, ct); _taskStore.FailTask(taskId, JsonElement.Parse("""{"code":-32000,"message":"something went wrong"}"""));
failedTask.SetResult(true);
}, ct);
};
await Assert.ThrowsAsync<McpException>(async () => await Assert.ThrowsAsync<McpException>(async () =>
await client.CallToolAsync( await client.CallToolAsync(
new CallToolRequestParams { Name = "async-tool" }, new CallToolRequestParams { Name = "async-tool" },
ct)); ct));
Assert.True(await failedTask.Task);
} }
[Fact] [Fact]
@ -198,17 +205,25 @@ public class McpServerTaskTests : ClientServerTestBase
await using var client = await CreateMcpClientForServer(); await using var client = await CreateMcpClientForServer();
var ct = TestContext.Current.CancellationToken; var ct = TestContext.Current.CancellationToken;
_ = Task.Run(async () => var cancelledTask = new TaskCompletionSource<bool>();
// Run cancellation task once the task from the tool call is created
_taskStore.OnTaskCreated += taskId =>
{ {
await Task.Delay(100, ct); Task.Run(async () =>
var taskId = _taskStore.GetAllTaskIds().Single(); {
_taskStore.CancelTask(taskId); await Task.Delay(100, ct);
}, ct); _taskStore.CancelTask(taskId);
cancelledTask.SetResult(true);
}, ct);
};
await Assert.ThrowsAsync<OperationCanceledException>(async () => await Assert.ThrowsAsync<OperationCanceledException>(async () =>
await client.CallToolAsync( await client.CallToolAsync(
new CallToolRequestParams { Name = "async-tool" }, new CallToolRequestParams { Name = "async-tool" },
ct)); ct));
Assert.True(await cancelledTask.Task);
} }
[Fact] [Fact]
@ -538,106 +553,135 @@ public class McpServerTaskTests : ClientServerTestBase
/// </summary> /// </summary>
private sealed class InMemoryTaskStore private sealed class InMemoryTaskStore
{ {
private readonly ConcurrentDictionary<string, TaskEntry> _tasks = new(); private readonly Dictionary<string, TaskEntry> _tasks = new();
internal Action<string>? OnTaskCreated;
public string CreateTask(McpTaskStatus initialStatus = McpTaskStatus.Working) public string CreateTask(McpTaskStatus initialStatus = McpTaskStatus.Working)
{ {
var taskId = Guid.NewGuid().ToString("N"); var taskId = Guid.NewGuid().ToString("N");
_tasks[taskId] = new TaskEntry lock (_tasks)
{ {
Status = initialStatus, _tasks[taskId] = new TaskEntry
CreatedAt = DateTimeOffset.UtcNow, {
LastUpdatedAt = DateTimeOffset.UtcNow, Status = initialStatus,
}; CreatedAt = DateTimeOffset.UtcNow,
LastUpdatedAt = DateTimeOffset.UtcNow,
};
}
OnTaskCreated?.Invoke(taskId);
return taskId; return taskId;
} }
public IEnumerable<string> GetAllTaskIds() => _tasks.Keys; public IEnumerable<string> GetAllTaskIds()
{
lock (_tasks)
{
return _tasks.Keys.ToArray();
}
}
public GetTaskResult GetTask(string taskId) public GetTaskResult GetTask(string taskId)
{ {
if (!_tasks.TryGetValue(taskId, out var entry)) lock (_tasks)
{ {
throw new McpException($"Unknown task: '{taskId}'"); if (!_tasks.TryGetValue(taskId, out var entry))
} {
throw new McpException($"Unknown task: '{taskId}'");
}
return entry.Status switch return entry.Status switch
{
McpTaskStatus.Working => new WorkingTaskResult
{ {
TaskId = taskId, McpTaskStatus.Working => new WorkingTaskResult
CreatedAt = entry.CreatedAt, {
LastUpdatedAt = entry.LastUpdatedAt, TaskId = taskId,
PollIntervalMs = 50, CreatedAt = entry.CreatedAt,
}, LastUpdatedAt = entry.LastUpdatedAt,
McpTaskStatus.Completed => new CompletedTaskResult PollIntervalMs = 50,
{ },
TaskId = taskId, McpTaskStatus.Completed => new CompletedTaskResult
CreatedAt = entry.CreatedAt, {
LastUpdatedAt = entry.LastUpdatedAt, TaskId = taskId,
Result = JsonSerializer.SerializeToElement(entry.Result, McpJsonUtilities.DefaultOptions), CreatedAt = entry.CreatedAt,
}, LastUpdatedAt = entry.LastUpdatedAt,
McpTaskStatus.Failed => new FailedTaskResult Result = JsonSerializer.SerializeToElement(entry.Result, McpJsonUtilities.DefaultOptions),
{ },
TaskId = taskId, McpTaskStatus.Failed => new FailedTaskResult
CreatedAt = entry.CreatedAt, {
LastUpdatedAt = entry.LastUpdatedAt, TaskId = taskId,
Error = entry.Error!.Value, CreatedAt = entry.CreatedAt,
}, LastUpdatedAt = entry.LastUpdatedAt,
McpTaskStatus.Cancelled => new CancelledTaskResult Error = entry.Error!.Value,
{ },
TaskId = taskId, McpTaskStatus.Cancelled => new CancelledTaskResult
CreatedAt = entry.CreatedAt, {
LastUpdatedAt = entry.LastUpdatedAt, TaskId = taskId,
}, CreatedAt = entry.CreatedAt,
McpTaskStatus.InputRequired => new InputRequiredTaskResult LastUpdatedAt = entry.LastUpdatedAt,
{ },
TaskId = taskId, McpTaskStatus.InputRequired => new InputRequiredTaskResult
CreatedAt = entry.CreatedAt, {
LastUpdatedAt = entry.LastUpdatedAt, TaskId = taskId,
InputRequests = entry.InputRequests ?? new Dictionary<string, InputRequest>(), CreatedAt = entry.CreatedAt,
}, LastUpdatedAt = entry.LastUpdatedAt,
_ => throw new InvalidOperationException($"Unexpected status: {entry.Status}") InputRequests = entry.InputRequests ?? new Dictionary<string, InputRequest>(),
}; },
_ => throw new InvalidOperationException($"Unexpected status: {entry.Status}")
};
}
} }
public void CompleteTask(string taskId, CallToolResult result) public void CompleteTask(string taskId, CallToolResult result)
{ {
if (_tasks.TryGetValue(taskId, out var entry)) lock (_tasks)
{ {
entry.Status = McpTaskStatus.Completed; if (_tasks.TryGetValue(taskId, out var entry))
entry.Result = result; {
entry.LastUpdatedAt = DateTimeOffset.UtcNow; entry.Result = result;
entry.LastUpdatedAt = DateTimeOffset.UtcNow;
entry.Status = McpTaskStatus.Completed;
}
} }
} }
public void FailTask(string taskId, JsonElement error) public void FailTask(string taskId, JsonElement error)
{ {
if (_tasks.TryGetValue(taskId, out var entry)) lock (_tasks)
{ {
entry.Status = McpTaskStatus.Failed; if (_tasks.TryGetValue(taskId, out var entry))
entry.Error = error; {
entry.LastUpdatedAt = DateTimeOffset.UtcNow; entry.Error = error;
entry.LastUpdatedAt = DateTimeOffset.UtcNow;
entry.Status = McpTaskStatus.Failed;
}
} }
} }
public void CancelTask(string taskId) public void CancelTask(string taskId)
{ {
if (_tasks.TryGetValue(taskId, out var entry)) lock (_tasks)
{ {
entry.Status = McpTaskStatus.Cancelled; if (_tasks.TryGetValue(taskId, out var entry))
entry.LastUpdatedAt = DateTimeOffset.UtcNow; {
entry.LastUpdatedAt = DateTimeOffset.UtcNow;
entry.Status = McpTaskStatus.Cancelled;
}
} }
} }
public void ProvideInput(string taskId, IDictionary<string, InputResponse> inputResponses) public void ProvideInput(string taskId, IDictionary<string, InputResponse> inputResponses)
{ {
if (_tasks.TryGetValue(taskId, out var entry)) lock (_tasks)
{ {
entry.InputResponses = inputResponses; if (_tasks.TryGetValue(taskId, out var entry))
// Transition back to working after receiving input {
entry.Status = McpTaskStatus.Working; entry.InputResponses = inputResponses;
entry.LastUpdatedAt = DateTimeOffset.UtcNow; entry.LastUpdatedAt = DateTimeOffset.UtcNow;
// Transition back to working after receiving input
entry.Status = McpTaskStatus.Working;
}
} }
} }

View file

@ -15,6 +15,8 @@ namespace ModelContextProtocol.Tests.Server;
/// </summary> /// </summary>
public class TaskPollStuckDetectorTests : ClientServerTestBase public class TaskPollStuckDetectorTests : ClientServerTestBase
{ {
private int _pollCount = 0;
public TaskPollStuckDetectorTests(ITestOutputHelper testOutputHelper) : base(testOutputHelper) public TaskPollStuckDetectorTests(ITestOutputHelper testOutputHelper) : base(testOutputHelper)
{ {
#if !NET #if !NET
@ -48,6 +50,8 @@ public class TaskPollStuckDetectorTests : ClientServerTestBase
// misbehaving-server condition the stuck-detector exists to break out of. // misbehaving-server condition the stuck-detector exists to break out of.
options.Handlers.GetTaskHandler = (context, cancellationToken) => options.Handlers.GetTaskHandler = (context, cancellationToken) =>
{ {
Interlocked.Increment(ref _pollCount);
return new ValueTask<GetTaskResult>(new InputRequiredTaskResult return new ValueTask<GetTaskResult>(new InputRequiredTaskResult
{ {
TaskId = context.Params!.TaskId, TaskId = context.Params!.TaskId,
@ -77,19 +81,13 @@ public class TaskPollStuckDetectorTests : ClientServerTestBase
await using var client = await CreateMcpClientForServer(); await using var client = await CreateMcpClientForServer();
var ct = TestContext.Current.CancellationToken; var ct = TestContext.Current.CancellationToken;
var sw = System.Diagnostics.Stopwatch.StartNew();
var ex = await Assert.ThrowsAsync<McpException>(async () => var ex = await Assert.ThrowsAsync<McpException>(async () =>
await client.CallToolAsync(new CallToolRequestParams { Name = "any-tool" }, ct)); await client.CallToolAsync(new CallToolRequestParams { Name = "any-tool" }, ct));
sw.Stop();
Assert.Contains(McpTaskStatus.InputRequired.ToString(), ex.Message); Assert.Contains(McpTaskStatus.InputRequired.ToString(), ex.Message);
Assert.Contains("consecutive polls", ex.Message); Assert.Contains("consecutive polls", ex.Message);
// 60 polls × 5ms ≈ 300ms; allow generous slack for CI. Assert.Equal(60, _pollCount);
Assert.True(sw.Elapsed < TimeSpan.FromSeconds(10),
$"Stuck-detector should give up promptly but took {sw.Elapsed}.");
} }
[Fact] [Fact]
@ -111,6 +109,7 @@ public class TaskPollStuckDetectorTests : ClientServerTestBase
// The message embeds the configured threshold, which is the strongest signal that the // The message embeds the configured threshold, which is the strongest signal that the
// option value (not the 60-default constant) is what governed the loop. // option value (not the 60-default constant) is what governed the loop.
Assert.Contains($"{CustomThreshold} consecutive polls", ex.Message); Assert.Contains($"{CustomThreshold} consecutive polls", ex.Message);
Assert.Equal(CustomThreshold, _pollCount);
} }
[Theory] [Theory]