Update to Orleans 8.2 and add support for RegisterGrainTimer (#163)

src/OrleansTestKit/OrleansTestKit.csproj

@@ -12,7 +12,7 @@
     <PackageTags>Orleans Cloud-Computing Actor-Model Actors Distributed-Systems C# .NET Test Testing</PackageTags>
     <PublishRepositoryUrl>true</PublishRepositoryUrl>
     <SymbolPackageFormat>snupkg</SymbolPackageFormat>
-    <Version>8.1.0</Version>
+    <Version>8.2.0</Version>
   </PropertyGroup>
 
   <PropertyGroup>
@@ -24,9 +24,9 @@
       <PrivateAssets>all</PrivateAssets>
       <IncludeAssets>runtime; build; native; contentfiles; analyzers; buildtransitive</IncludeAssets>
     </PackageReference>
-    <PackageReference Include="Microsoft.Orleans.Reminders" Version="8.1.0" />
-    <PackageReference Include="Microsoft.Orleans.Streaming" Version="8.1.0" />
-    <PackageReference Include="Microsoft.Orleans.Runtime" Version="8.1.0" />
+    <PackageReference Include="Microsoft.Orleans.Reminders" Version="8.2.0" />
+    <PackageReference Include="Microsoft.Orleans.Streaming" Version="8.2.0" />
+    <PackageReference Include="Microsoft.Orleans.Runtime" Version="8.2.0" />
     <PackageReference Include="Microsoft.SourceLink.GitHub" Version="1.1.1">
       <IncludeAssets>runtime; build; native; contentfiles; analyzers</IncludeAssets>
       <PrivateAssets>all</PrivateAssets>

src/OrleansTestKit/TestGrainActivationContext.cs

@@ -27,6 +27,8 @@ public sealed class TestGrainActivationContext : IGrainContext
     /// <inheritdoc/>
     public object GrainInstance { get; set; } = default!;
 
+    public void Migrate(Dictionary<string, object>? requestContext, CancellationToken cancellationToken = new CancellationToken()) => throw new NotImplementedException();
+
     /// <inheritdoc/>
     public GrainReference GrainReference { get; set; } = default!;
 
@@ -62,6 +64,10 @@ public sealed class TestGrainActivationContext : IGrainContext
     /// <inheritdoc/>
     public void ReceiveMessage(object message) => throw new NotImplementedException();
 
+    public void Activate(Dictionary<string, object>? requestContext, CancellationToken cancellationToken = new CancellationToken()) => throw new NotImplementedException();
+
+    public void Deactivate(DeactivationReason deactivationReason, CancellationToken cancellationToken = new CancellationToken()) => throw new NotImplementedException();
+
     /// <inheritdoc/>
     public void Rehydrate(IRehydrationContext context) => throw new NotImplementedException();
 

src/OrleansTestKit/TestGrainCreator.cs

@@ -36,7 +36,7 @@ public TestGrainCreator(IGrainRuntime runtime, IReminderRegistry reminderRegistr
         _runtime = runtime ?? throw new ArgumentNullException(nameof(runtime));
         _reminderRegistry = reminderRegistry;
         _serviceProvider = serviceProvider ?? throw new ArgumentNullException(nameof(runtime));
-        _contextProperty = typeof(Grain).GetProperty(GRAINCONTEXT_PROPERTYNAME, BindingFlags.Instance | BindingFlags.NonPublic);
+        _contextProperty = typeof(Grain).GetProperty(GRAINCONTEXT_PROPERTYNAME, BindingFlags.Instance | BindingFlags.Public);
         _runtimeProperty = typeof(Grain).GetProperty(RUNTIME_PROPERTYNAME, BindingFlags.Instance | BindingFlags.NonPublic);
         _contextPropertyBase = typeof(IGrainBase).GetProperty(GRAINCONTEXT_PROPERTYNAME, BindingFlags.Instance | BindingFlags.Public);
     }

src/OrleansTestKit/Timers/TestTimer.cs

@@ -1,20 +1,38 @@
 namespace Orleans.TestKit.Timers;
 
+/// <summary>
+/// Represents a specialization of TestTimer that is used for grain timers.
+/// </summary>
+public sealed class TestGrainTimer(Func<object?, CancellationToken, Task> asyncCallback, object state)
+    : TestTimer(asyncCallback, state), IGrainTimer
+{
+    public void Change(TimeSpan dueTime, TimeSpan period) => throw new NotSupportedException();
+}
+
 /// <summary>
 /// A test timer
 /// </summary>
-public sealed class TestTimer : IDisposable
+public class TestTimer : IDisposable
 {
     private Func<Task>? _asyncCallback;
+    public readonly CancellationTokenSource Cts = new();
 
     /// <summary>
     /// Initializes a new instance of the <see cref="TestTimer"/> class.
     /// </summary>
     /// <param name="asyncCallback">A callback function to invoke when the timer is fired</param>
     /// <param name="state">The timer's state</param>
-    public TestTimer(Func<object?, Task> asyncCallback, object? state) =>
+    public TestTimer(Func<object?, Task> asyncCallback, object state) =>
         _asyncCallback = () => asyncCallback(state);
 
+    /// <summary>
+    /// Initializes a new instance of the <see cref="TestTimer"/> class.
+    /// </summary>
+    /// <param name="asyncCallback">A callback function to invoke when the timer is fired</param>
+    /// <param name="state">The timer's state</param>
+    public TestTimer(Func<object?, CancellationToken, Task> asyncCallback, object state) =>
+        _asyncCallback = () => asyncCallback(state, Cts.Token);
+
     /// <summary>
     /// Gets a value indicating whether the timer has been disposed of
     /// </summary>
@@ -25,6 +43,10 @@ public void Dispose()
     {
         _asyncCallback = null;
         IsDisposed = true;
+
+        // used by TestGrainTimer callback only
+        Cts.Cancel();
+        Cts.Dispose();
     }
 
     /// <summary>
@@ -36,3 +58,4 @@ public Task FireAsync() =>
             ? throw new ObjectDisposedException(GetType().FullName)
             : _asyncCallback();
 }
+

src/OrleansTestKit/Timers/TestTimerRegistry.cs

@@ -1,5 +1,4 @@
 using Moq;
-using Orleans.Runtime;
 using Orleans.Timers;
 
 namespace Orleans.TestKit.Timers;
@@ -41,6 +40,7 @@ public async Task FireAllAsync()
     public Task FireAsync(int index) => _timers[index].FireAsync();
 
     /// <inheritdoc/>
+    [Obsolete]
     public IDisposable RegisterTimer(IGrainContext grainContext, Func<object?, Task> asyncCallback, object? state, TimeSpan dueTime, TimeSpan period)
     {
         if (grainContext == null)
@@ -49,8 +49,28 @@ public IDisposable RegisterTimer(IGrainContext grainContext, Func<object?, Task>
         }
 
         Mock.Object.RegisterTimer(grainContext, asyncCallback, state, dueTime, period);
-        var timer = new TestTimer(asyncCallback, state);
+
+        var timer = new TestTimer(asyncCallback, state!);
         _timers.Add(timer);
+
         return timer;
     }
+
+    public IGrainTimer RegisterGrainTimer<TState>(IGrainContext grainContext, Func<TState, CancellationToken, Task> asyncCallback, TState state,
+        GrainTimerCreationOptions options)
+    {
+        if (grainContext == null)
+        {
+            throw new ArgumentNullException(nameof(grainContext));
+        }
+
+        var cb = new Func<object?, CancellationToken, Task>((s, ct) => asyncCallback(((TState?)s)!, ct));
+
+        Mock.Object.RegisterGrainTimer(grainContext, asyncCallback, state, options);
+
+        var grainTimer = new TestGrainTimer(cb, state!);
+        _timers.Add(grainTimer);
+
+        return grainTimer;
+    }
 }

test/OrleansTestKit.Tests/Grains/HelloTimers.cs

@@ -2,6 +2,10 @@
 
 public class HelloTimers : Grain<HelloTimersState>, IGrainWithIntegerKey
 {
+    internal const int GrainTimer0 = 3;
+
+    internal IGrainTimer _grainTimer0 = default!;
+
     private IDisposable? _secretTimer;
 
     private IDisposable _timer0 = default!;
@@ -16,6 +20,8 @@ public override Task OnActivateAsync(CancellationToken cancellationToken)
         _timer1 = RegisterTimer(_ => OnTimer1(), null, TimeSpan.Zero, TimeSpan.FromSeconds(1));
         _timer2 = RegisterTimer(_ => OnTimer2(), null, TimeSpan.Zero, TimeSpan.FromSeconds(1));
 
+        _grainTimer0 = this.RegisterGrainTimer<object?>((_, c) => OnGrainTimer0(c), null, TimeSpan.Zero, TimeSpan.FromSeconds(1));
+
         return base.OnActivateAsync(cancellationToken);
     }
 
@@ -40,6 +46,21 @@ private Task OnSecretTimer()
         return Task.CompletedTask;
     }
 
+    private async Task OnGrainTimer0(CancellationToken cancellationToken)
+    {
+        State.GrainTimer0Fired = true;
+
+        var delay = Task.Delay(1000, cancellationToken);
+        try
+        {
+            await delay;
+        }
+        catch (TaskCanceledException)
+        {
+            State.GrainTimer0Cancelled = true;
+        }
+    }
+
     private Task OnTimer0()
     {
         State.Timer0Fired = true;
@@ -74,4 +95,8 @@ public HelloTimersState()
     public bool Timer1Fired { get; set; }
 
     public bool Timer2Fired { get; set; }
+
+    public bool GrainTimer0Fired { get; set; }
+
+    public bool GrainTimer0Cancelled { get; set; }
 }

test/OrleansTestKit.Tests/OrleansTestKit.Tests.csproj

@@ -14,7 +14,7 @@
     <PackageReference Include="FluentAssertions" Version="6.12.0" />
     <PackageReference Include="Microsoft.CodeCoverage" Version="17.8.0" />
     <PackageReference Include="Microsoft.NET.Test.Sdk" Version="17.8.0" />
-    <PackageReference Include="Microsoft.Orleans.Sdk" Version="8.1.0" />
+    <PackageReference Include="Microsoft.Orleans.Sdk" Version="8.2.0" />
     <PackageReference Include="NSubstitute" Version="5.1.0" />
     <PackageReference Include="System.Configuration.ConfigurationManager" Version="7.0.0" />
     <PackageReference Include="xunit" Version="2.6.5" />

test/OrleansTestKit.Tests/Tests/TimerTests.cs

@@ -1,11 +1,45 @@
 using FluentAssertions;
+using Orleans.TestKit.Timers;
 using TestGrains;
 using Xunit;
 
 namespace Orleans.TestKit.Tests;
 
 public class TimerTests : TestKitBase
 {
+    [Fact]
+    public async Task ShouldFirstGrainTimerAsync()
+    {
+        // Arrange
+        var grain = await Silo.CreateGrainAsync<HelloTimers>(0);
+
+        // Act
+        await Silo.FireTimerAsync(HelloTimers.GrainTimer0);
+
+        // Assert
+        var state = Silo.State<HelloTimers, HelloTimersState>();
+        state.GrainTimer0Fired.Should().BeTrue();
+        state.GrainTimer0Cancelled.Should().BeFalse();
+    }
+
+    [Fact]
+        public async Task ShouldCancelFirstGrainTimerAsync()
+    {
+        // Arrange
+        var grain = await Silo.CreateGrainAsync<HelloTimers>(0);
+
+        // Act
+        _ = Silo.FireTimerAsync(HelloTimers.GrainTimer0);
+        await Task.Delay(100);
+        grain._grainTimer0.Dispose();
+        await Task.Delay(100);
+
+        // Assert
+        var state = Silo.State<HelloTimers, HelloTimersState>();
+        state.GrainTimer0Fired.Should().BeTrue();
+        state.GrainTimer0Cancelled.Should().BeTrue();
+    }
+
     [Fact]
     public async Task ShouldFireAllTimersAsync()
     {
@@ -19,6 +53,8 @@ public async Task ShouldFireAllTimersAsync()
         var state = Silo.State<HelloTimers, HelloTimersState>();
         state.Timer0Fired.Should().BeTrue();
         state.Timer1Fired.Should().BeTrue();
+        state.Timer2Fired.Should().BeTrue();
+        state.GrainTimer0Fired.Should().BeTrue();
     }
 
     [Fact]
@@ -36,6 +72,7 @@ public async Task ShouldFireAllTimersRepeatedlyAsync()
         state.Timer0Fired.Should().BeTrue();
         state.Timer1Fired.Should().BeTrue();
         state.Timer2Fired.Should().BeTrue();
+        state.GrainTimer0Fired.Should().BeTrue();
 
     }