From 19f9ab1d3cfeb5156436b89fd4c15403bc3eb2d3 Mon Sep 17 00:00:00 2001
From: Justin Baur <19896123+justindbaur@users.noreply.github.com>
Date: Wed, 4 Dec 2024 13:32:13 -0500
Subject: [PATCH] Use ParameterResolver

---
 .../DatabaseDataAttribute.cs                  |  2 +-
 .../Utilities/AutoMigrateAttribute.cs         |  5 +++--
 .../Utilities/ServiceProviderTheoryDataRow.cs | 19 +++++++++++--------
 .../Utilities/TestCustomizerAttribute.cs      |  4 ++--
 4 files changed, 17 insertions(+), 13 deletions(-)

diff --git a/test/Infrastructure.IntegrationTest/DatabaseDataAttribute.cs b/test/Infrastructure.IntegrationTest/DatabaseDataAttribute.cs
index 2460234f54..d9a44e3b90 100644
--- a/test/Infrastructure.IntegrationTest/DatabaseDataAttribute.cs
+++ b/test/Infrastructure.IntegrationTest/DatabaseDataAttribute.cs
@@ -43,7 +43,7 @@ public class DatabaseDataAttribute : DataAttribute
             }
             else
             {
-                theory = new ServiceTheoryDataRow(testMethod, disposalTracker, customizationContext.Services.BuildServiceProvider());
+                theory = new ServiceTheoryDataRow(testMethod, disposalTracker, customizationContext);
             }
 
             theory
diff --git a/test/Infrastructure.IntegrationTest/Utilities/AutoMigrateAttribute.cs b/test/Infrastructure.IntegrationTest/Utilities/AutoMigrateAttribute.cs
index e4d6a406ef..0e791aa964 100644
--- a/test/Infrastructure.IntegrationTest/Utilities/AutoMigrateAttribute.cs
+++ b/test/Infrastructure.IntegrationTest/Utilities/AutoMigrateAttribute.cs
@@ -1,6 +1,5 @@
 
 using Bit.Core.Enums;
-using Bit.Core.Utilities;
 using Bit.Infrastructure.IntegrationTest.Services;
 using Microsoft.Extensions.DependencyInjection;
 
@@ -32,6 +31,8 @@ public class AutoMigrateAttribute : TestCustomizerAttribute
         // Build services provider early and run migrations
         var sp = customizationContext.Services.BuildServiceProvider();
         var migrator = sp.GetRequiredService<IMigrationTesterService>();
-        migrator.ApplyMigration()
+        migrator.ApplyMigration();
+
+        return Task.CompletedTask;
     }
 }
diff --git a/test/Infrastructure.IntegrationTest/Utilities/ServiceProviderTheoryDataRow.cs b/test/Infrastructure.IntegrationTest/Utilities/ServiceProviderTheoryDataRow.cs
index 9db04a9f07..1e758e847e 100644
--- a/test/Infrastructure.IntegrationTest/Utilities/ServiceProviderTheoryDataRow.cs
+++ b/test/Infrastructure.IntegrationTest/Utilities/ServiceProviderTheoryDataRow.cs
@@ -2,7 +2,6 @@ using System.Reflection;
 using Microsoft.Extensions.DependencyInjection;
 using Xunit;
 using Xunit.Sdk;
-using Xunit.v3;
 
 namespace Bit.Infrastructure.IntegrationTest.Utilities;
 
@@ -10,25 +9,29 @@ public class ServiceTheoryDataRow : TheoryDataRowBase
 {
     private readonly MethodInfo _testMethod;
     private readonly DisposalTracker _disposalTracker;
-    private readonly IServiceProvider _serviceProvider;
+    private readonly CustomizationContext _customizationContext;
 
-    public ServiceTheoryDataRow(MethodInfo testMethod, DisposalTracker disposalTracker, IServiceProvider serviceProvider)
+    public ServiceTheoryDataRow(MethodInfo testMethod, DisposalTracker disposalTracker, CustomizationContext customizationContext)
     {
         _testMethod = testMethod;
         _disposalTracker = disposalTracker;
-        _serviceProvider = serviceProvider;
+        _customizationContext = customizationContext;
     }
 
     protected override object?[] GetData()
     {
         var parameters = _testMethod.GetParameters();
-        // Create a scope for this test
-        var scope = _serviceProvider.CreateAsyncScope();
 
-        var objects = new object[parameters.Length];
+        var sp = _customizationContext.Services.BuildServiceProvider();
+
+        // Create a scope for this test
+        var scope = sp.CreateAsyncScope();
+
+        var objects = new object?[parameters.Length];
         for (var i = 0; i < parameters.Length; i++)
         {
-            objects[i] = scope.ServiceProvider.GetRequiredService(parameters[i].ParameterType);
+            var parameter = parameters[i];
+            objects[i] = _customizationContext.ParameterResolver(scope.ServiceProvider, parameter);
         }
 
         _disposalTracker.Add(scope);
diff --git a/test/Infrastructure.IntegrationTest/Utilities/TestCustomizerAttribute.cs b/test/Infrastructure.IntegrationTest/Utilities/TestCustomizerAttribute.cs
index acd086d109..758ace1387 100644
--- a/test/Infrastructure.IntegrationTest/Utilities/TestCustomizerAttribute.cs
+++ b/test/Infrastructure.IntegrationTest/Utilities/TestCustomizerAttribute.cs
@@ -16,7 +16,7 @@ public class CustomizationContext
 
     public IServiceCollection Services { get; }
 
-    public Func<ServiceProvider, ParameterInfo, object?> ParameterResolver { get; set; } = DefaultParameterResolver;
+    public Func<IServiceProvider, ParameterInfo, object?> ParameterResolver { get; set; } = DefaultParameterResolver;
 
 
     public CustomizationContext(Database database, MethodInfo testMethod, DisposalTracker disposalTracker)
@@ -27,7 +27,7 @@ public class CustomizationContext
         Services = new ServiceCollection();
     }
 
-    private static object? DefaultParameterResolver(ServiceProvider services, ParameterInfo parameter)
+    private static object? DefaultParameterResolver(IServiceProvider services, ParameterInfo parameter)
     {
         return services.GetService(parameter.ParameterType);
     }