From d0dd776157961a882e181b8bdae57f28cbaf734d Mon Sep 17 00:00:00 2001
From: jrmccannon <jmccannon@bitwarden.com>
Date: Wed, 20 Nov 2024 11:49:24 -0600
Subject: [PATCH] Fixing up the tests.

---
 .../SingleOrgPolicyValidator.cs               |  5 +-
 .../SingleOrgPolicyValidatorTests.cs          | 51 +++++++++++++++----
 2 files changed, 42 insertions(+), 14 deletions(-)

diff --git a/src/Core/AdminConsole/OrganizationFeatures/Policies/PolicyValidators/SingleOrgPolicyValidator.cs b/src/Core/AdminConsole/OrganizationFeatures/Policies/PolicyValidators/SingleOrgPolicyValidator.cs
index 683171012..249b71b72 100644
--- a/src/Core/AdminConsole/OrganizationFeatures/Policies/PolicyValidators/SingleOrgPolicyValidator.cs
+++ b/src/Core/AdminConsole/OrganizationFeatures/Policies/PolicyValidators/SingleOrgPolicyValidator.cs
@@ -128,11 +128,10 @@ public class SingleOrgPolicyValidator : IPolicyValidator
             ou.Type != OrganizationUserType.Owner &&
             ou.Type != OrganizationUserType.Admin &&
             ou.UserId != savingUserId
-            ).ToList();
+        ).ToList();
 
         var userOrgs = await _organizationUserRepository.GetManyByManyUsersAsync(
-                removableOrgUsers.Select(ou => ou.UserId!.Value));
-
+            removableOrgUsers.Select(ou => ou.UserId!.Value));
         foreach (var orgUser in removableOrgUsers)
         {
             if (userOrgs.Any(ou => ou.UserId == orgUser.UserId
diff --git a/test/Core.Test/AdminConsole/OrganizationFeatures/Policies/PolicyValidators/SingleOrgPolicyValidatorTests.cs b/test/Core.Test/AdminConsole/OrganizationFeatures/Policies/PolicyValidators/SingleOrgPolicyValidatorTests.cs
index d5e09694b..83d10a7ec 100644
--- a/test/Core.Test/AdminConsole/OrganizationFeatures/Policies/PolicyValidators/SingleOrgPolicyValidatorTests.cs
+++ b/test/Core.Test/AdminConsole/OrganizationFeatures/Policies/PolicyValidators/SingleOrgPolicyValidatorTests.cs
@@ -11,6 +11,7 @@ using Bit.Core.Auth.Repositories;
 using Bit.Core.Context;
 using Bit.Core.Entities;
 using Bit.Core.Enums;
+using Bit.Core.Models.Commands;
 using Bit.Core.Models.Data.Organizations.OrganizationUsers;
 using Bit.Core.Repositories;
 using Bit.Core.Services;
@@ -62,8 +63,11 @@ public class SingleOrgPolicyValidatorTests
         Assert.True(string.IsNullOrEmpty(result));
     }
 
-    [Theory, BitAutoData]
+    [Theory]
+    [BitAutoData(true)]
+    [BitAutoData(false)]
     public async Task OnSaveSideEffectsAsync_RemovesNonCompliantUsers(
+        bool isAccountDeprovisioningEnabled,
         [PolicyUpdate(PolicyType.SingleOrg)] PolicyUpdate policyUpdate,
         [Policy(PolicyType.SingleOrg, false)] Policy policy,
         Guid savingUserId,
@@ -117,19 +121,40 @@ public class SingleOrgPolicyValidatorTests
         sutProvider.GetDependency<ICurrentContext>().UserId.Returns(savingUserId);
         sutProvider.GetDependency<IOrganizationRepository>().GetByIdAsync(policyUpdate.OrganizationId).Returns(organization);
 
+        sutProvider.GetDependency<IFeatureService>().IsEnabled(FeatureFlagKeys.AccountDeprovisioning)
+            .Returns(isAccountDeprovisioningEnabled);
+
+        sutProvider.GetDependency<IRevokeNonCompliantOrganizationUserCommand>()
+            .RevokeNonCompliantOrganizationUsersAsync(Arg.Any<RevokeOrganizationUsersRequest>())
+            .Returns(new CommandResult());
+
         await sutProvider.Sut.OnSaveSideEffectsAsync(policyUpdate, policy);
 
-        await sutProvider.GetDependency<IRemoveOrganizationUserCommand>()
-            .Received(1)
-            .RemoveUserAsync(policyUpdate.OrganizationId, nonCompliantUser.Id, savingUserId);
-        await sutProvider.GetDependency<IMailService>()
-            .Received(1)
-            .SendOrganizationUserRemovedForPolicySingleOrgEmailAsync(organization.DisplayName(),
-                "user3@example.com");
+        if (isAccountDeprovisioningEnabled)
+        {
+            await sutProvider.GetDependency<IRevokeNonCompliantOrganizationUserCommand>()
+                .Received(1)
+                .RevokeNonCompliantOrganizationUsersAsync(Arg.Any<RevokeOrganizationUsersRequest>());
+            await sutProvider.GetDependency<IMailService>()
+                .Received(1)
+                .SendOrganizationUserRevokedForPolicySingleOrgEmailAsync(organization.DisplayName(),
+                    "user3@example.com");
+        }
+        else
+        {
+            await sutProvider.GetDependency<IRemoveOrganizationUserCommand>()
+                .Received(1)
+                .RemoveUserAsync(policyUpdate.OrganizationId, nonCompliantUser.Id, savingUserId);
+            await sutProvider.GetDependency<IMailService>()
+                .Received(1)
+                .SendOrganizationUserRemovedForPolicySingleOrgEmailAsync(organization.DisplayName(),
+                    "user3@example.com");
+        }
+
     }
 
     [Theory, BitAutoData]
-    public async Task OnSaveSideEffectsAsync_GivenUpdatedSingleOrgPolicy_WhenAccountDeprovisioningIsDisabled_Then(
+    public async Task OnSaveSideEffectsAsync_WhenAccountDeprovisioningIsEnabled_ThenUsersAreRevoked(
         [PolicyUpdate(PolicyType.SingleOrg)] PolicyUpdate policyUpdate,
         [Policy(PolicyType.SingleOrg, false)] Policy policy,
         Guid savingUserId,
@@ -184,12 +209,16 @@ public class SingleOrgPolicyValidatorTests
         sutProvider.GetDependency<IOrganizationRepository>().GetByIdAsync(policyUpdate.OrganizationId)
             .Returns(organization);
 
-        sutProvider.GetDependency<IFeatureService>().IsEnabled(FeatureFlagKeys.AccountDeprovisioning).Returns(false);
+        sutProvider.GetDependency<IFeatureService>().IsEnabled(FeatureFlagKeys.AccountDeprovisioning).Returns(true);
+
+        sutProvider.GetDependency<IRevokeNonCompliantOrganizationUserCommand>()
+            .RevokeNonCompliantOrganizationUsersAsync(Arg.Any<RevokeOrganizationUsersRequest>())
+            .Returns(new CommandResult());
 
         await sutProvider.Sut.OnSaveSideEffectsAsync(policyUpdate, policy);
 
         await sutProvider.GetDependency<IRevokeNonCompliantOrganizationUserCommand>()
-            .DidNotReceive()
+            .Received()
             .RevokeNonCompliantOrganizationUsersAsync(Arg.Any<RevokeOrganizationUsersRequest>());
     }
 }