Use TimeProvider in InMemoryMcpTaskStore for deterministic testing (#1303)
This commit is contained in:
parent
5da0fe506f
commit
b967835144
2 changed files with 88 additions and 22 deletions
|
|
@ -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
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
|
|
|||
Loading…
Reference in a new issue