TestHost - Client ability to use IHostConfigurator to align w/ Silo functionality

src/Orleans.TestingHost/TestClusterBuilder.cs

@@ -8,9 +8,7 @@
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Hosting;
 using Orleans.Hosting;
-using Orleans.Internal;
 using Orleans.Runtime;
-using Orleans.Runtime.TestHooks;
 
 namespace Orleans.TestingHost
 {
@@ -80,7 +78,7 @@ public Func<string, IConfiguration, Task<SiloHandle>> CreateSiloAsync
                 Options.ConnectionTransport = ConnectionTransportType.TcpSocket;
             }
         }
-        
+
         /// <summary>
         /// Adds a configuration delegate to the builder
         /// </summary>        
@@ -121,13 +119,18 @@ public TestClusterBuilder ConfigureHostConfiguration(Action<IConfigurationBuilde
         }
 
         /// <summary>
-        /// Adds the client builder configurator, which must implement <see cref="IClientBuilderConfigurator"/> or <see cref="IHostConfigurator"/>.
+        /// Adds an implementation of <see cref="IClientBuilderConfigurator"/> or <see cref="IHostConfigurator"/> to configure the client created for the test cluster
         /// </summary>
-        /// <typeparam name="TClientBuilderConfigurator">The client builder type</typeparam>
+        /// <typeparam name="T">The client builder type</typeparam>
         /// <returns>The builder.</returns>
-        public TestClusterBuilder AddClientBuilderConfigurator<TClientBuilderConfigurator>() where TClientBuilderConfigurator : IClientBuilderConfigurator, new()
+        public TestClusterBuilder AddClientBuilderConfigurator<T>() where T : new()
         {
-            this.Options.ClientBuilderConfiguratorTypes.Add(typeof(TClientBuilderConfigurator).AssemblyQualifiedName);
+            if (!typeof(IClientBuilderConfigurator).IsAssignableFrom(typeof(T)) && !typeof(IHostConfigurator).IsAssignableFrom(typeof(T)))
+            {
+                throw new ArgumentException($"The type {typeof(T)} is not assignable to either {nameof(IClientBuilderConfigurator)} or {nameof(IHostConfigurator)}");
+            }
+
+            this.Options.ClientBuilderConfiguratorTypes.Add(typeof(T).AssemblyQualifiedName);
             return this;
         }
 
@@ -155,7 +158,7 @@ public TestCluster Build()
             var configuration = configBuilder.Build();
             var finalOptions = new TestClusterOptions();
             configuration.Bind(finalOptions);
-            
+
             var configSources = new ReadOnlyCollection<IConfigurationSource>(configBuilder.Sources);
             var testCluster = new TestCluster(finalOptions, configSources, portAllocator);
             if (CreateSiloAsync != null) testCluster.CreateSiloAsync = CreateSiloAsync;

src/Orleans.TestingHost/TestClusterHostFactory.cs

@@ -3,22 +3,20 @@
 using System.Diagnostics;
 using System.Linq;
 using System.Net;
-using Microsoft.Extensions.Hosting;
 using Microsoft.Extensions.Configuration;
 using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Hosting;
 using Newtonsoft.Json;
-using Orleans.Hosting;
-using Orleans.Runtime.TestHooks;
 using Orleans.Configuration;
+using Orleans.Configuration.Internal;
+using Orleans.Hosting;
 using Orleans.Messaging;
 using Orleans.Runtime;
 using Orleans.Runtime.MembershipService;
+using Orleans.Runtime.TestHooks;
 using Orleans.Statistics;
-using Orleans.TestingHost.Utils;
 using Orleans.TestingHost.Logging;
-using Orleans.Configuration.Internal;
-using Microsoft.Extensions.Hosting.Internal;
-using Orleans.TestingHost.InMemoryTransport;
+using Orleans.TestingHost.Utils;
 
 namespace Orleans.TestingHost
 {
@@ -88,12 +86,16 @@ public static IHost CreateClusterClient(string hostName, IConfiguration configur
             var hostBuilder = new HostBuilder();
             hostBuilder.UseEnvironment(Environments.Development);
             hostBuilder.Properties["Configuration"] = configuration;
-            hostBuilder.ConfigureHostConfiguration(cb => cb.AddConfiguration(configuration))
-                .UseOrleansClient((ctx, clientBuilder) =>
-                {
-                    clientBuilder.Configure<ClusterOptions>(configuration);
-                    ConfigureAppServices(configuration, clientBuilder);
-                })
+            hostBuilder.ConfigureHostConfiguration(cb => cb.AddConfiguration(configuration));
+
+            hostBuilder.UseOrleansClient(clientBuilder =>
+            {
+                clientBuilder.Configure<ClusterOptions>(configuration);
+            });
+
+            ConfigureClientAppServices(configuration, hostBuilder);
+
+            hostBuilder
                 .ConfigureServices(services =>
                 {
                     TryConfigureClientMembership(configuration, services);
@@ -176,7 +178,7 @@ private static void ConfigureAppServices(IConfiguration configuration, IHostBuil
             }
         }
 
-        private static void ConfigureAppServices(IConfiguration configuration, IClientBuilder clientBuilder)
+        private static void ConfigureClientAppServices(IConfiguration configuration, IHostBuilder hostBuilder)
         {
             var builderConfiguratorTypes = configuration.GetSection(nameof(TestClusterOptions.ClientBuilderConfiguratorTypes))?.Get<string[]>();
             if (builderConfiguratorTypes == null) return;
@@ -185,8 +187,14 @@ private static void ConfigureAppServices(IConfiguration configuration, IClientBu
             {
                 if (!string.IsNullOrWhiteSpace(builderConfiguratorType))
                 {
-                    var builderConfigurator = (IClientBuilderConfigurator)Activator.CreateInstance(Type.GetType(builderConfiguratorType, true));
-                    builderConfigurator?.Configure(configuration, clientBuilder);
+                    var builderConfigurator = Activator.CreateInstance(Type.GetType(builderConfiguratorType, true));
+
+                    (builderConfigurator as IHostConfigurator)?.Configure(hostBuilder);
+
+                    if (builderConfigurator is IClientBuilderConfigurator clientBuilderConfigurator)
+                    {
+                        hostBuilder.UseOrleansClient(clientBuilder => clientBuilderConfigurator.Configure(configuration, clientBuilder));
+                    }
                 }
             }
         }

test/TestInfrastructure/Orleans.TestingHost.Tests/TestClusterTests.cs

@@ -1,6 +1,6 @@
-using System;
-using System.Threading.Tasks;
-using Orleans.Hosting;
+using FluentAssertions;
+using Microsoft.Extensions.Configuration;
+using Microsoft.Extensions.Hosting;
 using Orleans.TestingHost.Tests.Grains;
 using TestExtensions;
 using Xunit;
@@ -155,6 +155,58 @@ public async Task CanInitialize()
         }
     }
 
+    public class T10
+    {
+        private static bool _hostWasInvoked;
+        private static bool _clientWasInvoked;
+
+        [Fact, TestCategory("Functional")]
+        public async Task ClientBuilder_HostConfigurator() => await Test<HostConfigurator>(true, false);
+
+        [Fact, TestCategory("Functional")]
+        public async Task ClientBuilder_ClientHostConfigurator() => await Test<ClientHostConfigurator>(true, true);
+
+        [Fact, TestCategory("Functional")]
+        public async Task ClientBuilder_ClientConfigurator() => await Test<ClientConfigurator>(false, true);
+
+        private static async Task Test<TConfigurator>(bool hostInvoked, bool clientInvoked)
+            where TConfigurator : new()
+        {
+            _hostWasInvoked = false;
+            _clientWasInvoked = false;
+
+            var builder = new TestClusterBuilder(2);
+            builder.Options.ServiceId = Guid.NewGuid().ToString();
+            builder.AddClientBuilderConfigurator<TConfigurator>();
+            using var testCluster = builder.Build();
+
+            await testCluster.DeployAsync();
+            await testCluster.StopAllSilosAsync();
+
+            _hostWasInvoked.Should().Be(hostInvoked);
+            _clientWasInvoked.Should().Be(clientInvoked);
+        }
+
+        private class HostConfigurator : IHostConfigurator
+        {
+            public void Configure(IHostBuilder hostBuilder) => _hostWasInvoked = true;
+        }
+
+        private class ClientHostConfigurator : IHostConfigurator, IClientBuilderConfigurator
+        {
+
+            public void Configure(IHostBuilder hostBuilder) => _hostWasInvoked = true;
+            public void Configure(IConfiguration configuration, IClientBuilder clientBuilder) => _clientWasInvoked = true;
+        }
+
+        private class ClientConfigurator : IClientBuilderConfigurator
+        {
+            public void Configure(IConfiguration configuration, IClientBuilder clientBuilder) => _clientWasInvoked = true;
+        }
+    }
+
+
+
     public class TestClusterTests : IDisposable, IAsyncLifetime
     {
         private readonly ITestOutputHelper output;