From 05192a9a8dffa29c78213d47227df94cf78ddf9c Mon Sep 17 00:00:00 2001
From: Thomas Rittson <trittson@bitwarden.com>
Date: Tue, 10 Oct 2023 14:15:43 +1000
Subject: [PATCH] Add EnableFeatureFlag extension method for tests

---
 test/Common/AutoFixture/SutProviderExtensions.cs  | 12 ++++++++++++
 test/Core.Test/Services/CollectionServiceTests.cs |  6 +-----
 2 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/test/Common/AutoFixture/SutProviderExtensions.cs b/test/Common/AutoFixture/SutProviderExtensions.cs
index 1fdf226539..e3a45e9b95 100644
--- a/test/Common/AutoFixture/SutProviderExtensions.cs
+++ b/test/Common/AutoFixture/SutProviderExtensions.cs
@@ -1,4 +1,6 @@
 using AutoFixture;
+using Bit.Core;
+using Bit.Core.Context;
 using Bit.Core.Services;
 using Bit.Core.Settings;
 using NSubstitute;
@@ -47,4 +49,14 @@ public static class SutProviderExtensions
             .SetDependency(mockHttpClientFactory)
             .Create();
     }
+
+    public static SutProvider<T> EnableFeatureFlag<T>(this SutProvider<T> sutProvider,
+        string featureFlagKey)
+    {
+        sutProvider.GetDependency<IFeatureService>()
+            .IsEnabled(featureFlagKey, Arg.Any<ICurrentContext>(), Arg.Any<bool>())
+            .Returns(true);
+
+        return sutProvider;
+    }
 }
diff --git a/test/Core.Test/Services/CollectionServiceTests.cs b/test/Core.Test/Services/CollectionServiceTests.cs
index 08a0f437fa..ab40b574ef 100644
--- a/test/Core.Test/Services/CollectionServiceTests.cs
+++ b/test/Core.Test/Services/CollectionServiceTests.cs
@@ -112,11 +112,7 @@ public class CollectionServiceTest
     {
         collection.Id = default;
         sutProvider.GetDependency<IOrganizationRepository>().GetByIdAsync(organization.Id).Returns(organization);
-
-        // TODO: move to customization / attribute
-        sutProvider.GetDependency<IFeatureService>()
-            .IsEnabled(FeatureFlagKeys.FlexibleCollections, Arg.Any<ICurrentContext>())
-            .Returns(true);
+        sutProvider.EnableFeatureFlag(FeatureFlagKeys.FlexibleCollections);
 
         var ex = await Assert.ThrowsAsync<BadRequestException>(() => sutProvider.Sut.SaveAsync(collection, null, users));
         Assert.Contains("At least one member or group must have can manage permission.", ex.Message);