Test new way of capturing the service provider

Author: andreasohlund

src/NServiceBus.Storage.MongoDB.TransactionalSession.AcceptanceTests/Infrastructure/CaptureServiceProviderStartupTask.cs

@@ -3,18 +3,11 @@
     using System;
     using System.Threading;
     using System.Threading.Tasks;
-    using NServiceBus.AcceptanceTesting;
-    using NServiceBus.Features;
+    using Features;
 
     public class CaptureServiceProviderStartupTask : FeatureStartupTask
     {
-        public CaptureServiceProviderStartupTask(IServiceProvider serviceProvider, ScenarioContext context)
-        {
-            if (context is IInjectServiceProvider c)
-            {
-                c.ServiceProvider = serviceProvider;
-            }
-        }
+        public CaptureServiceProviderStartupTask(IServiceProvider serviceProvider, TransactionalSessionTestContext context) => context.RegisterServiceProvider(serviceProvider);
 
         protected override Task OnStart(IMessageSession session, CancellationToken cancellationToken = default) => Task.CompletedTask;
 

src/NServiceBus.Storage.MongoDB.TransactionalSession.AcceptanceTests/Infrastructure/DefaultServer.cs

@@ -35,9 +35,9 @@ public virtual async Task<EndpointConfiguration> GetConfiguration(RunDescriptor
 
         endpointConfiguration.GetSettings().Set(persistence);
 
-        if (!typeof(IDoNotCaptureServiceProvider).IsAssignableFrom(endpointCustomizations.BuilderType))
+        if (runDescriptor.ScenarioContext is TransactionalSessionTestContext testContext)
         {
-            endpointConfiguration.RegisterStartupTask(sp => new CaptureServiceProviderStartupTask(sp, runDescriptor.ScenarioContext));
+            endpointConfiguration.RegisterStartupTask(sp => new CaptureServiceProviderStartupTask(sp, testContext));
         }
 
         await configurationBuilderCustomization(endpointConfiguration).ConfigureAwait(false);

src/NServiceBus.Storage.MongoDB.TransactionalSession.AcceptanceTests/Infrastructure/IDoNotCaptureServiceProvider.cs

@@ -0,0 +1,39 @@
+namespace NServiceBus.TransactionalSession.AcceptanceTests;
+
+using System;
+using System.Collections.Concurrent;
+using AcceptanceTesting;
+
+public class TransactionalSessionTestContext : ScenarioContext
+{
+    public IServiceProvider ServiceProvider
+    {
+        get
+        {
+            var endpointName = GetCurrentEndpoint();
+
+            if (!serviceProviders.TryGetValue(endpointName, out var serviceProvider))
+            {
+                throw new InvalidOperationException("Could not find service provider for endpoint " + endpointName);
+            }
+
+            return serviceProvider;
+        }
+    }
+
+    public void RegisterServiceProvider(IServiceProvider serviceProvider)
+    {
+        var endpointName = GetCurrentEndpoint();
+
+        serviceProviders[endpointName] = serviceProvider;
+    }
+
+    string GetCurrentEndpoint()
+    {
+        var endpointName = GetType().GetProperty("CurrentEndpoint")!.GetValue(this, null) as string;
+
+        return endpointName;
+    }
+
+    readonly ConcurrentDictionary<string, IServiceProvider> serviceProviders = new();
+}

src/NServiceBus.Storage.MongoDB.TransactionalSession.AcceptanceTests/Infrastructure/TransactionalSessionTestContext.cs

@@ -37,11 +37,9 @@ public async Task Should_send_messages_on_transactional_session_commit()
         Assert.That(context.MessageReceived, Is.True);
     }
 
-    class Context : ScenarioContext, IInjectServiceProvider
+    class Context : TransactionalSessionTestContext
     {
         public bool MessageReceived { get; set; }
-
-        public IServiceProvider ServiceProvider { get; set; }
     }
 
     class SendOnlyEndpoint : EndpointConfigurationBuilder

src/NServiceBus.Storage.MongoDB.TransactionalSession.AcceptanceTests/When_using_outbox_send_only.cs

@@ -1,7 +1,5 @@
 namespace NServiceBus.TransactionalSession.AcceptanceTests
 {
-    using System;
-    using System.Threading;
     using System.Threading.Tasks;
     using AcceptanceTesting;
     using Microsoft.Extensions.DependencyInjection;
@@ -166,12 +164,11 @@ await mongoSession.Client.GetDatabase(SetupFixture.DatabaseName)
             Assert.That(documents.ToList().Count, Is.EqualTo(1));
         }
 
-        class Context : ScenarioContext, IInjectServiceProvider
+        class Context : TransactionalSessionTestContext
         {
             public bool MessageReceived { get; set; }
             public bool CompleteMessageReceived { get; set; }
             public string SessionId { get; set; }
-            public IServiceProvider ServiceProvider { get; set; }
         }
 
         class AnEndpoint : EndpointConfigurationBuilder