Unit Testing Grain State in Orleans.

This post will show how to unit test the state of your grains when using the Microsoft Orleans Testing Host and assumes you are already familiar with setting up the Testing Host. If you are not, I recommend first reading the official documentation.

My Requirements

I have a relatively simple test host setup for Orleans which follows the documentation linked above and therefore simplifies the implementation of storage testing.

My setup is as follows:

• A single TestCluster fixture is re-used for all tests.
• A single PersistentState is injected into each grain.
• Each code path within the grain that mutates state calls WriteStateAsync().

Faking the Grain Storage

We will first implement an IGrainStorage provider to store all of the silo state in a concurrent dictionary for thread-safe access. The key will be the grain reference, and the value will be the grain’s state.

public class FakeGrainStorage : IGrainStorage
{
    public ConcurrentDictionary<GrainReference, IGrainState> Storage { get; } = new();

    public Task ClearStateAsync(string grainType, GrainReference grainReference, IGrainState grainState) 
        => Task.FromResult(Storage.TryRemove(grainReference, out _));

    public Task ReadStateAsync(string grainType, GrainReference grainReference, IGrainState grainState) 
        => Task.FromResult(Storage.TryGetValue(grainReference, out grainState));

    public Task WriteStateAsync(string grainType, GrainReference grainReference, IGrainState grainState) 
        => Task.FromResult(Storage.TryAdd(grainReference, grainState));
}

We can then register this with our silo builder.

public class ClusterFixture : IAsyncLifetime
{
    ....
    public static FakeGrainStorage GrainStorage { get; } = new();

    class TestSiloConfiguration : ISiloConfigurator
    {
        public void Configure(ISiloBuilder siloBuilder)
            => siloBuilder
                ....
                .ConfigureServices(services =>
                {
                    services
                        .AddSingleton<IGrainStorage>(GrainStorage)
                        ....
                });
    }
}

Accessing the Grain Storage

Now that we have registered our grain storage with the silos, we need to access it. To do this, we can add a convenience method to the FakeGrainStorage class to help access the state for a specific grain.

public TState GetGrainState<TState>(IGrain grain)
    => Storage.TryGetValue((GrainReference)grain, out var state)
        ? (TState)state.State
        : default;

An Example Test

Now that we have set up our storage, let’s look at how to test against it:

[Fact]
public async Task ExampleTest()
{
    // Arrange
    var myGrain = _cluster.GrainFactory.GetGrain<IMyGrain>(Guid.NewGuid());

    // Act
    await myGrain.MethodThatMutatesAndSavesStateAsync();

    // Assert
    var myGrainState = ClusterFixture.GrainStorage.GetGrainState<MyGrainState>(myGrain);
    myGrainState.SomeProperty.Should().Be("Test");
}

Note that your grain must call WriteStateAsync() before the assertion, for the state to be added to the fake storage. My grains persist their mutated state at the end of every code path, so this works well, but you will need to alter this testing approach if you call WriteStateAsync() at a different point, such as on a timer or silo shutdown.

Summary

In this post, I demonstrated how to register simple in-memory grain storage that can be used to execute unit test assertions.

It would be an improvement to the Orleans TestingHost framework if TestClusterBuilder.AddSiloBuilderConfigurator had an overload that allowed passing in an instance of an ISiloConfigurator rather than just a type, as this would allow us to remove the static nature of the FakeGrainStorage held in the ClusterFixture` which will cause issues for anyone using multiple clusters in parallel tests.

If you do use multiple clusters in parallel, or you are injecting multiple PersistentState into a single grain constructor, I suggest taking the above approach and adapting it to your needs, using this sample from the Orleans unit testing project as a guide.