From 955579d10bc05ede5abb549b7ac45cac2361260d Mon Sep 17 00:00:00 2001
From: Thomas Rittson <trittson@bitwarden.com>
Date: Fri, 18 Oct 2024 13:33:58 +1000
Subject: [PATCH] Use extension method for FakeTimeProvider

---
 test/Common/AutoFixture/SutProvider.cs            | 10 ----------
 test/Common/AutoFixture/SutProviderExtensions.cs  | 15 +++++++++++++++
 .../Policies/SavePolicyCommandTests.cs            |  5 ++---
 3 files changed, 17 insertions(+), 13 deletions(-)

diff --git a/test/Common/AutoFixture/SutProvider.cs b/test/Common/AutoFixture/SutProvider.cs
index 0e8287cde7d5..5bf563f88c82 100644
--- a/test/Common/AutoFixture/SutProvider.cs
+++ b/test/Common/AutoFixture/SutProvider.cs
@@ -128,16 +128,6 @@ public object Create(object request, ISpecimenContext context)
                 return _sutProvider.GetDependency(parameterInfo.ParameterType, "");
             }
 
-            // TimeProvider should use FakeTimeProvider rather than a Substitute mock
-            if (parameterInfo.ParameterType == typeof(TimeProvider))
-            {
-                var fakeTimeProvider = new FakeTimeProvider();
-                _sutProvider.SetDependency(parameterInfo.ParameterType, fakeTimeProvider, parameterInfo.Name);
-                // Also register it under FakeTimeProvider so it can be easily retrieved in GetDependency
-                _sutProvider.SetDependency(typeof(FakeTimeProvider), fakeTimeProvider);
-                return fakeTimeProvider;
-            }
-
             // This is the equivalent of _fixture.Create<parameterInfo.ParameterType>, but no overload for
             // Create(Type type) exists.
             var dependency = new SpecimenContext(_fixture).Resolve(new SeededRequest(parameterInfo.ParameterType,
diff --git a/test/Common/AutoFixture/SutProviderExtensions.cs b/test/Common/AutoFixture/SutProviderExtensions.cs
index 1fdf22653943..f40c19b3bd5b 100644
--- a/test/Common/AutoFixture/SutProviderExtensions.cs
+++ b/test/Common/AutoFixture/SutProviderExtensions.cs
@@ -1,6 +1,7 @@
 using AutoFixture;
 using Bit.Core.Services;
 using Bit.Core.Settings;
+using Microsoft.Extensions.Time.Testing;
 using NSubstitute;
 using RichardSzalay.MockHttp;
 
@@ -47,4 +48,18 @@ public static SutProvider<T> ConfigureBaseIdentityClientService<T>(this SutProvi
             .SetDependency(mockHttpClientFactory)
             .Create();
     }
+
+    /// <summary>
+    /// Configures SutProvider to use FakeTimeProvider.
+    /// It is registered under both the TimeProvider type and the FakeTimeProvider type
+    /// so that it can be retrieved in a type-safe manner with GetDependency.
+    /// </summary>
+    public static SutProvider<T> ConfigureFakeTimeProvider<T>(this SutProvider<T> sutProvider)
+    {
+        var fakeTimeProvider = new FakeTimeProvider();
+        return sutProvider
+            .SetDependency((TimeProvider)fakeTimeProvider)
+            .SetDependency(fakeTimeProvider)
+            .Create();
+    }
 }
diff --git a/test/Core.Test/AdminConsole/OrganizationFeatures/Policies/SavePolicyCommandTests.cs b/test/Core.Test/AdminConsole/OrganizationFeatures/Policies/SavePolicyCommandTests.cs
index d09d9f2d4b1b..c1b9640a27de 100644
--- a/test/Core.Test/AdminConsole/OrganizationFeatures/Policies/SavePolicyCommandTests.cs
+++ b/test/Core.Test/AdminConsole/OrganizationFeatures/Policies/SavePolicyCommandTests.cs
@@ -264,9 +264,8 @@ private static SutProvider<SavePolicyCommand> SutProviderFactory(IEnumerable<IPo
     {
         var fixture = new Fixture();
         fixture.Customizations.Add(new PolicyValidatorSpecimenBuilder(policyValidators ?? new List<IPolicyValidator>()));
-        var sutProvider = new SutProvider<SavePolicyCommand>(fixture);
-        sutProvider.Create();
-        return sutProvider;
+        return new SutProvider<SavePolicyCommand>(fixture)
+            .ConfigureFakeTimeProvider();
     }
 
     private static void ArrangeOrganization(SutProvider<SavePolicyCommand> sutProvider, PolicyUpdate policyUpdate)