Fix keyed services key type

src/OrleansTestKit/Services/TestServiceProvider.cs

@@ -5,106 +5,87 @@
 
 namespace Orleans.TestKit.Services;
 
-/// <summary>
-/// The test service provider
-/// </summary>
+/// <summary>The test service provider</summary>
 public sealed class TestServiceProvider : IServiceProvider, IKeyedServiceProvider
 {
+    private readonly Dictionary<(object?, Type), object> _keyedServices = new();
+
     private readonly TestKitOptions _options;
+
     private readonly Dictionary<Type, object> _services = new();
-    private readonly Dictionary<(object?, Type), object> _keyedServices = new();
 
-    /// <summary>
-    /// Initializes a new instance of the <see cref="TestServiceProvider"/> class.
-    /// </summary>
+    /// <summary>Initializes a new instance of the <see cref="TestServiceProvider"/> class.</summary>
     /// <param name="options">The test kit options to use</param>
     public TestServiceProvider(TestKitOptions options)
     {
         ArgumentNullException.ThrowIfNull(options);
         _options = options;
     }
 
-    /// <summary>
-    /// Adds or updates a service to the provider
-    /// </summary>
+    /// <summary>Adds or updates a keyed service to the provider</summary>
     /// <typeparam name="T">The service type</typeparam>
+    /// <param name="serviceKey">The service key</param>
     /// <param name="instance">The instance to add</param>
     /// <returns>The instance</returns>
     /// <exception cref="ArgumentNullException">Instance must be not null</exception>
-    public T AddService<T>(T instance)
+    public T AddKeyedService<T>(object? serviceKey, T instance)
     {
         ArgumentNullException.ThrowIfNull(instance);
 
-        _services[typeof(T)] = instance;
+        _keyedServices[(serviceKey, typeof(T))] = instance;
         return instance;
     }
 
-    /// <summary>
-    /// Adds or updates a keyed service to the provider
-    /// </summary>
+    /// <summary>Adds or updates a keyed mock to the service provider</summary>
+    /// <typeparam name="T">The underlying service type</typeparam>
+    /// <param name="serviceKey">The service key</param>
+    /// <returns>The newly created mock</returns>
+    public Mock<T> AddKeyedServiceProbe<T>(object? serviceKey)
+        where T : class =>
+        AddKeyedServiceProbe(serviceKey, new Mock<T>());
+
+    /// <summary>Adds or updates a keyed mock to the service provider</summary>
+    /// <typeparam name="T">The underlying service type</typeparam>
+    /// <param name="serviceKey">The service key</param>
+    /// <param name="mock">The mock to add</param>
+    /// <returns>The mock</returns>
+    public Mock<T> AddKeyedServiceProbe<T>(object? serviceKey, Mock<T> mock)
+        where T : class
+    {
+        AddKeyedService(serviceKey, mock.Object);
+        return mock;
+    }
+
+    /// <summary>Adds or updates a service to the provider</summary>
     /// <typeparam name="T">The service type</typeparam>
-    /// <param name="name">The service key</param>
     /// <param name="instance">The instance to add</param>
     /// <returns>The instance</returns>
     /// <exception cref="ArgumentNullException">Instance must be not null</exception>
-    public T AddKeyedService<T>(string name, T instance)
+    public T AddService<T>(T instance)
     {
         ArgumentNullException.ThrowIfNull(instance);
 
-        _keyedServices[(name, typeof(T))] = instance;
+        _services[typeof(T)] = instance;
         return instance;
     }
 
-    /// <summary>
-    /// Adds a mock to the service provider
-    /// </summary>
+    /// <summary>Adds or updates a mock to the service provider</summary>
     /// <typeparam name="T">The underlying service type</typeparam>
     /// <param name="mock">The mock to add</param>
     /// <returns>The mock</returns>
-    public Mock<T> AddServiceProbe<T>(Mock<T> mock) where T : class
+    public Mock<T> AddServiceProbe<T>(Mock<T> mock)
+        where T : class
     {
         AddService(mock.Object);
         return mock;
     }
 
-    /// <summary>
-    /// Adds a mock to the service provider
-    /// </summary>
+    /// <summary>Adds a mock to the service provider</summary>
     /// <typeparam name="T">The underlying service type</typeparam>
     /// <returns>The newly created mock</returns>
-    public Mock<T> AddServiceProbe<T>() where T : class => AddServiceProbe(new Mock<T>());
-
-    /// <inheritdoc/>
-    public object GetService(Type serviceType)
-    {
-        ArgumentNullException.ThrowIfNull(serviceType);
-
-        if (_services.TryGetValue(serviceType, out var service))
-        {
-            return service;
-        }
-
-        // If using strict service probes, throw the exception
-        if (_options.StrictServiceProbes)
-        {
-            throw new Exception($"Service probe does not exist for type {serviceType.Name}. Ensure that it is added before the grain is tested.");
-        }
-        else
-        {
-            // Create a new mock
-            if (Activator.CreateInstance(typeof(Mock<>).MakeGenericType(serviceType)) is not IMock<object> mock)
-            {
-                throw new Exception($"Failed to instantiate {serviceType.Name}.");
-            }
-
-            service = mock.Object;
-
-            // Save the newly created grain for the next call
-            _services.Add(serviceType, service);
-
-            return service;
-        }
-    }
+    public Mock<T> AddServiceProbe<T>()
+        where T : class =>
+        AddServiceProbe(new Mock<T>());
 
     public object? GetKeyedService(Type serviceType, object? serviceKey)
     {
@@ -155,4 +136,36 @@ public object GetRequiredKeyedService(Type serviceType, object? serviceKey)
 
         return service;
     }
+
+    /// <inheritdoc/>
+    public object GetService(Type serviceType)
+    {
+        ArgumentNullException.ThrowIfNull(serviceType);
+
+        if (_services.TryGetValue(serviceType, out var service))
+        {
+            return service;
+        }
+
+        // If using strict service probes, throw the exception
+        if (_options.StrictServiceProbes)
+        {
+            throw new Exception($"Service probe does not exist for type {serviceType.Name}. Ensure that it is added before the grain is tested.");
+        }
+        else
+        {
+            // Create a new mock
+            if (Activator.CreateInstance(typeof(Mock<>).MakeGenericType(serviceType)) is not IMock<object> mock)
+            {
+                throw new Exception($"Failed to instantiate {serviceType.Name}.");
+            }
+
+            service = mock.Object;
+
+            // Save the newly created grain for the next call
+            _services.Add(serviceType, service);
+
+            return service;
+        }
+    }
 }

src/OrleansTestKit/Services/TestServicesExtensions.cs

@@ -4,14 +4,47 @@ namespace Orleans.TestKit;
 
 public static class TestServicesExtensions
 {
-    public static T AddService<T>(this TestKitSilo silo, T instance)
+    public static T AddKeyedService<T>(this TestKitSilo silo, object? serviceKey, T instance)
+        where T : class
+    {
+        if (silo == null)
+        {
+            throw new ArgumentNullException(nameof(silo));
+        }
+
+        return silo.ServiceProvider.AddKeyedService(serviceKey, instance);
+    }
+
+    public static Mock<T> AddKeyedServiceProbe<T>(this TestKitSilo silo, object? serviceKey)
+        where T : class
+    {
+        if (silo == null)
+        {
+            throw new ArgumentNullException(nameof(silo));
+        }
+
+        return silo.ServiceProvider.AddKeyedServiceProbe<T>(serviceKey);
+    }
+
+    public static Mock<T> AddKeyedServiceProbe<T>(this TestKitSilo silo, object? serviceKey, Mock<T> mock)
         where T : class
     {
         if (silo == null)
         {
             throw new ArgumentNullException(nameof(silo));
         }
 
+        return silo.ServiceProvider.AddKeyedServiceProbe(serviceKey, mock);
+    }
+
+    public static T AddService<T>(this TestKitSilo silo, T instance)
+            where T : class
+    {
+        if (silo == null)
+        {
+            throw new ArgumentNullException(nameof(silo));
+        }
+
         return silo.ServiceProvider.AddService(instance);
     }
 

test/OrleansTestKit.Tests/Grains/DependencyGrain.cs

@@ -0,0 +1,37 @@
+using Microsoft.Extensions.DependencyInjection;
+using TestInterfaces;
+
+namespace TestGrains;
+
+public interface IDependency
+{
+    string? GetValue();
+}
+
+public sealed class DependencyGrain : Grain, IDependencyGrain
+{
+    private readonly IDependency? _firstKeyedDependency;
+
+    private readonly IDependency? _secondKeyedDependency;
+
+    private readonly IDependency? _unkeyedDependency;
+
+    public DependencyGrain(
+        IDependency? unkeyedDependency,
+        [FromKeyedServices("first")] IDependency? firstKeyedDependency,
+        [FromKeyedServices("second")] IDependency? secondKeyedDependency)
+    {
+        _unkeyedDependency = unkeyedDependency;
+        _firstKeyedDependency = firstKeyedDependency;
+        _secondKeyedDependency = secondKeyedDependency;
+    }
+
+    public Task<string?> GetFirstKeyedServiceValue() =>
+        Task.FromResult(_firstKeyedDependency?.GetValue());
+
+    public Task<string?> GetSecondKeyedServiceValue() =>
+        Task.FromResult(_secondKeyedDependency?.GetValue());
+
+    public Task<string?> GetUnkeyedServiceValue() =>
+        Task.FromResult(_unkeyedDependency?.GetValue());
+}

test/OrleansTestKit.Tests/Interfaces/IDependencyGrain.cs

@@ -0,0 +1,10 @@
+namespace TestInterfaces;
+
+public interface IDependencyGrain : IGrainWithGuidKey
+{
+    Task<string?> GetFirstKeyedServiceValue();
+
+    Task<string?> GetSecondKeyedServiceValue();
+
+    Task<string?> GetUnkeyedServiceValue();
+}

test/OrleansTestKit.Tests/Tests/DependencyGrainTests.cs

@@ -0,0 +1,60 @@
+using FluentAssertions;
+using Moq;
+using TestGrains;
+using Xunit;
+
+namespace Orleans.TestKit.Tests;
+
+public class DependencyGrainTests : TestKitBase
+{
+    [Fact]
+    public async Task GrainWithoutServices()
+    {
+        // Arrange
+        var grain = await Silo.CreateGrainAsync<DependencyGrain>(Guid.NewGuid());
+
+        // Act
+        var unkeyedServiceValue = await grain.GetUnkeyedServiceValue();
+        var firstKeyedServiceValue = await grain.GetFirstKeyedServiceValue();
+        var secondKeyedServiceValue = await grain.GetSecondKeyedServiceValue();
+
+        // Assert
+        unkeyedServiceValue.Should().BeNull();
+        firstKeyedServiceValue.Should().BeNull();
+        secondKeyedServiceValue.Should().BeNull();
+    }
+
+    [Fact]
+    public async Task GrainWithServices()
+    {
+        // Arrange
+        var unkeyedService = new Mock<IDependency>(MockBehavior.Strict);
+        unkeyedService.Setup(x => x.GetValue()).Returns("");
+        Silo.ServiceProvider.AddServiceProbe(unkeyedService);
+
+        var firstKeyedService = new Mock<IDependency>(MockBehavior.Strict);
+        firstKeyedService.Setup(x => x.GetValue()).Returns("first");
+        Silo.ServiceProvider.AddKeyedServiceProbe("first", firstKeyedService);
+
+        var secondKeyedService = new Mock<IDependency>(MockBehavior.Strict);
+        secondKeyedService.Setup(x => x.GetValue()).Returns("second");
+        Silo.ServiceProvider.AddKeyedServiceProbe("second", secondKeyedService);
+
+        var grain = await Silo.CreateGrainAsync<DependencyGrain>(Guid.NewGuid());
+
+        // Act
+        var unkeyedServiceValue = await grain.GetUnkeyedServiceValue();
+        var firstKeyedServiceValue = await grain.GetFirstKeyedServiceValue();
+        var secondKeyedServiceValue = await grain.GetSecondKeyedServiceValue();
+
+        // Assert
+        unkeyedServiceValue.Should().BeEmpty();
+        unkeyedService.Verify(x => x.GetValue(), Times.Once);
+
+        firstKeyedServiceValue.Should().Be("first");
+        firstKeyedService.Verify(x => x.GetValue(), Times.Once);
+
+        secondKeyedServiceValue.Should().Be("second");
+        secondKeyedService.Verify(x => x.GetValue(), Times.Once);
+    }
+}