Use TimeProvider in InMemoryMcpTaskStore for deterministic testing (#1303)

This commit is contained in:
Copilot 2026-02-18 21:47:31 -08:00 • committed by GitHub
commit b967835144
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 88 additions and 22 deletions

View file

@ -35,7 +35,11 @@ public sealed class InMemoryMcpTaskStore : IMcpTaskStore, IDisposable
private readonly TimeSpan? _defaultTtl;
private readonly TimeSpan? _maxTtl;
private readonly TimeSpan _pollInterval;
#if MCP_TEST_TIME_PROVIDER
private readonly ITimer? _cleanupTimer;
#else
private readonly Timer? _cleanupTimer;
#endif
private readonly int _pageSize;
private readonly int? _maxTasks;
private readonly int? _maxTasksPerSession;
@ -134,7 +138,11 @@ public sealed class InMemoryMcpTaskStore : IMcpTaskStore, IDisposable
cleanupInterval ??= TimeSpan.FromMinutes(1);
if (cleanupInterval.Value != Timeout.InfiniteTimeSpan)
{
#if MCP_TEST_TIME_PROVIDER
_cleanupTimer = _timeProvider.CreateTimer(CleanupExpiredTasks, null, cleanupInterval.Value, cleanupInterval.Value);
#else
_cleanupTimer = new Timer(CleanupExpiredTasks, null, cleanupInterval.Value, cleanupInterval.Value);
#endif
}
}

View file

@ -235,14 +235,24 @@ public class InMemoryMcpTaskStoreTests : LoggedTest
[Fact]
public async Task UpdateTaskStatusAsync_UpdatesLastUpdatedAt()
{
// Arrange
using var store = new InMemoryMcpTaskStore();
// Arrange - Use FakeTimeProvider for deterministic testing
var fakeTime = new FakeTimeProvider(DateTimeOffset.UtcNow);
using var store = new TestInMemoryMcpTaskStore(
defaultTtl: null,
maxTtl: null,
pollInterval: null,
cleanupInterval: Timeout.InfiniteTimeSpan,
pageSize: 100,
maxTasks: null,
maxTasksPerSession: null,
timeProvider: fakeTime);
var metadata = new McpTaskMetadata();
var task = await store.CreateTaskAsync(metadata, new RequestId("req-1"), new JsonRpcRequest { Method = "test" }, null, TestContext.Current.CancellationToken);
var originalTimestamp = task.LastUpdatedAt;
// Wait a bit to ensure timestamp changes
await Task.Delay(10, TestContext.Current.CancellationToken);
// Advance time to ensure timestamp changes
fakeTime.Advance(TimeSpan.FromMilliseconds(10));
// Act
await store.UpdateTaskStatusAsync(task.TaskId, McpTaskStatus.Working, null, null, TestContext.Current.CancellationToken);
@ -460,16 +470,28 @@ public class InMemoryMcpTaskStoreTests : LoggedTest
[Fact]
public async Task Dispose_StopsCleanupTimer()
{
// Arrange
var store = new InMemoryMcpTaskStore(cleanupInterval: TimeSpan.FromMilliseconds(100));
var metadata = new McpTaskMetadata { TimeToLive = TimeSpan.FromMilliseconds(100) }; // Very short TTL
// Arrange - Use FakeTimeProvider for deterministic testing
var fakeTime = new FakeTimeProvider(DateTimeOffset.UtcNow);
var cleanupInterval = TimeSpan.FromMilliseconds(100);
var store = new TestInMemoryMcpTaskStore(
defaultTtl: null,
maxTtl: null,
pollInterval: null,
cleanupInterval: cleanupInterval,
pageSize: 100,
maxTasks: null,
maxTasksPerSession: null,
timeProvider: fakeTime);
var metadata = new McpTaskMetadata { TimeToLive = TimeSpan.FromMilliseconds(100) };
await store.CreateTaskAsync(metadata, new RequestId("req-1"), new JsonRpcRequest { Method = "test" }, null, TestContext.Current.CancellationToken);
// Act
store.Dispose();
// Wait longer than cleanup interval
await Task.Delay(300, TestContext.Current.CancellationToken);
// Advance time - timer should not fire after dispose
fakeTime.Advance(TimeSpan.FromTicks(cleanupInterval.Ticks * 3));
// Assert - Store should still be accessible after dispose (no exceptions)
// The cleanup timer should have stopped
@ -479,17 +501,33 @@ public class InMemoryMcpTaskStoreTests : LoggedTest
[Fact]
public async Task CleanupExpiredTasks_RemovesExpiredTasks()
{
// Arrange
using var store = new InMemoryMcpTaskStore(cleanupInterval: TimeSpan.FromMilliseconds(50));
var metadata = new McpTaskMetadata { TimeToLive = TimeSpan.FromMilliseconds(100) }; // 100ms TTL
// Arrange - Use FakeTimeProvider for deterministic testing
var fakeTime = new FakeTimeProvider(DateTimeOffset.UtcNow);
var cleanupInterval = TimeSpan.FromMilliseconds(50);
var ttl = TimeSpan.FromMilliseconds(100);
using var store = new TestInMemoryMcpTaskStore(
defaultTtl: null,
maxTtl: null,
pollInterval: null,
cleanupInterval: cleanupInterval,
pageSize: 100,
maxTasks: null,
maxTasksPerSession: null,
timeProvider: fakeTime);
var metadata = new McpTaskMetadata { TimeToLive = ttl };
var task = await store.CreateTaskAsync(metadata, new RequestId("req-1"), new JsonRpcRequest { Method = "test" }, null, TestContext.Current.CancellationToken);
// Verify task exists initially
var resultBefore = await store.ListTasksAsync(cancellationToken: TestContext.Current.CancellationToken);
Assert.Single(resultBefore.Tasks);
// Wait for task to expire and cleanup timer to run (wait for at least 3 cleanup cycles)
await Task.Delay(250, TestContext.Current.CancellationToken);
// Advance time past the TTL to make task expired
fakeTime.Advance(ttl + TimeSpan.FromMilliseconds(1));
// Trigger cleanup by advancing time past cleanup interval
fakeTime.Advance(cleanupInterval);
// Act - List tasks to verify cleanup happened
var resultAfter = await store.ListTasksAsync(cancellationToken: TestContext.Current.CancellationToken);
@ -787,8 +825,18 @@ public class InMemoryMcpTaskStoreTests : LoggedTest
[Fact]
public async Task ListTasksAsync_ConsistentWithExpiredTasksRemovedBetweenPages()
{
// Arrange - Use TTL of 1 second
using var store = new InMemoryMcpTaskStore(defaultTtl: TimeSpan.FromSeconds(1), pageSize: 5, cleanupInterval: Timeout.InfiniteTimeSpan);
// Arrange - Use FakeTimeProvider for deterministic testing
var fakeTime = new FakeTimeProvider(DateTimeOffset.UtcNow);
var ttl = TimeSpan.FromSeconds(1);
using var store = new TestInMemoryMcpTaskStore(
defaultTtl: ttl,
maxTtl: null,
pollInterval: null,
cleanupInterval: Timeout.InfiniteTimeSpan,
pageSize: 5,
maxTasks: null,
maxTasksPerSession: null,
timeProvider: fakeTime);
// Create 15 tasks
for (int i = 0; i < 15; i++)
@ -799,8 +847,8 @@ public class InMemoryMcpTaskStoreTests : LoggedTest
// Act - Get first page immediately
var result1 = await store.ListTasksAsync(cancellationToken: TestContext.Current.CancellationToken);
// Wait for tasks to expire
await Task.Delay(TimeSpan.FromSeconds(1.5), TestContext.Current.CancellationToken);
// Advance time past TTL to make tasks expire
fakeTime.Advance(ttl + TimeSpan.FromMilliseconds(500));
// Get second page after expiration
var result2 = await store.ListTasksAsync(cursor: result1.NextCursor, cancellationToken: TestContext.Current.CancellationToken);
@ -1016,16 +1064,26 @@ public class InMemoryMcpTaskStoreTests : LoggedTest
[Fact]
public async Task CreateTaskAsync_MaxTasksPerSession_ExcludesExpiredTasks()
{
// Arrange - Short TTL and per-session limit of 1
// Arrange - Use FakeTimeProvider for deterministic testing
var fakeTime = new FakeTimeProvider(DateTimeOffset.UtcNow);
var shortTtl = TimeSpan.FromMilliseconds(50);
using var store = new InMemoryMcpTaskStore(defaultTtl: shortTtl, maxTasksPerSession: 1);
using var store = new TestInMemoryMcpTaskStore(
defaultTtl: shortTtl,
maxTtl: null,
pollInterval: null,
cleanupInterval: Timeout.InfiniteTimeSpan,
pageSize: 100,
maxTasks: null,
maxTasksPerSession: 1,
timeProvider: fakeTime);
var metadata = new McpTaskMetadata();
// Create first task
await store.CreateTaskAsync(metadata, new RequestId("req-1"), new JsonRpcRequest { Method = "test" }, "session-1", TestContext.Current.CancellationToken);
// Wait for it to expire
await Task.Delay(100, TestContext.Current.CancellationToken);
// Advance time past TTL to make the first task expire
fakeTime.Advance(shortTtl + TimeSpan.FromMilliseconds(1));
// Should be able to create another task since the first one expired
var task2 = await store.CreateTaskAsync(metadata, new RequestId("req-2"), new JsonRpcRequest { Method = "test" }, "session-1", TestContext.Current.CancellationToken);