[PM-15547] Fix two-factor authentication revocation logic and update …

…related tests (#5246)

* Fix two-factor authentication revocation logic and update related tests

* Refine test for RevokeNonCompliantOrganizationUserCommand to assert single user revocation

src/Core/Services/Implementations/UserService.cs

@@ -1372,17 +1372,16 @@ public void SetTwoFactorProvider(User user, TwoFactorProviderType type, bool set
     private async Task CheckPoliciesOnTwoFactorRemovalAsync(User user)
     {
         var twoFactorPolicies = await _policyService.GetPoliciesApplicableToUserAsync(user.Id, PolicyType.TwoFactorAuthentication);
-        var organizationsManagingUser = await GetOrganizationsManagingUserAsync(user.Id);
 
         var removeOrgUserTasks = twoFactorPolicies.Select(async p =>
         {
             var organization = await _organizationRepository.GetByIdAsync(p.OrganizationId);
-            if (_featureService.IsEnabled(FeatureFlagKeys.AccountDeprovisioning) && organizationsManagingUser.Any(o => o.Id == p.OrganizationId))
+            if (_featureService.IsEnabled(FeatureFlagKeys.AccountDeprovisioning))
             {
                 await _revokeNonCompliantOrganizationUserCommand.RevokeNonCompliantOrganizationUsersAsync(
                     new RevokeOrganizationUsersRequest(
                         p.OrganizationId,
-                        [new OrganizationUserUserDetails { UserId = user.Id, OrganizationId = p.OrganizationId }],
+                        [new OrganizationUserUserDetails { Id = p.OrganizationUserId, OrganizationId = p.OrganizationId }],
                         new SystemUser(EventSystemUser.TwoFactorDisabled)));
                 await _mailService.SendOrganizationUserRevokedForTwoFactoryPolicyEmailAsync(organization.DisplayName(), user.Email);
             }

test/Core.Test/AdminConsole/OrganizationFeatures/OrganizationUsers/RevokeNonCompliantOrganizationUserCommandTests.cs

@@ -162,7 +162,7 @@ public async Task RevokeNonCompliantOrganizationUsersAsync_GivenValidPopulatedRe
 
         await sutProvider.GetDependency<IOrganizationUserRepository>()
             .Received(1)
-            .RevokeManyByIdAsync(Arg.Any<IEnumerable<Guid>>());
+            .RevokeManyByIdAsync(Arg.Is<IEnumerable<Guid>>(x => x.Count() == 1 && x.Contains(userToRevoke.Id)));
 
         Assert.True(result.Success);
 

test/Core.Test/Services/UserServiceTests.cs

@@ -454,8 +454,10 @@ await sutProvider.GetDependency<IMailService>()
     }
 
     [Theory, BitAutoData]
-    public async Task DisableTwoFactorProviderAsync_WithAccountDeprovisioningEnabled_WhenOrganizationHas2FAPolicyEnabled_WhenUserIsManaged_DisablingAllProviders_RemovesOrRevokesUserAndSendsEmail(
-        SutProvider<UserService> sutProvider, User user, Organization organization1, Organization organization2)
+    public async Task DisableTwoFactorProviderAsync_WithAccountDeprovisioningEnabled_WhenOrganizationHas2FAPolicyEnabled_DisablingAllProviders_RevokesUserAndSendsEmail(
+        SutProvider<UserService> sutProvider, User user,
+        Organization organization1, Guid organizationUserId1,
+        Organization organization2, Guid organizationUserId2)
     {
         // Arrange
         user.SetTwoFactorProviders(new Dictionary<TwoFactorProviderType, TwoFactorProvider>
@@ -464,6 +466,7 @@ public async Task DisableTwoFactorProviderAsync_WithAccountDeprovisioningEnabled
         });
         organization1.Enabled = organization2.Enabled = true;
         organization1.UseSso = organization2.UseSso = true;
+
         sutProvider.GetDependency<IFeatureService>()
             .IsEnabled(FeatureFlagKeys.AccountDeprovisioning)
             .Returns(true);
@@ -474,12 +477,14 @@ public async Task DisableTwoFactorProviderAsync_WithAccountDeprovisioningEnabled
                 new OrganizationUserPolicyDetails
                 {
                     OrganizationId = organization1.Id,
+                    OrganizationUserId = organizationUserId1,
                     PolicyType = PolicyType.TwoFactorAuthentication,
                     PolicyEnabled = true
                 },
                 new OrganizationUserPolicyDetails
                 {
                     OrganizationId = organization2.Id,
+                    OrganizationUserId = organizationUserId2,
                     PolicyType = PolicyType.TwoFactorAuthentication,
                     PolicyEnabled = true
                 }
@@ -490,9 +495,6 @@ public async Task DisableTwoFactorProviderAsync_WithAccountDeprovisioningEnabled
         sutProvider.GetDependency<IOrganizationRepository>()
             .GetByIdAsync(organization2.Id)
             .Returns(organization2);
-        sutProvider.GetDependency<IOrganizationRepository>()
-            .GetByVerifiedUserEmailDomainAsync(user.Id)
-            .Returns(new[] { organization1 });
         var expectedSavedProviders = JsonHelpers.LegacySerialize(new Dictionary<TwoFactorProviderType, TwoFactorProvider>(), JsonHelpers.LegacyEnumKeyResolver);
 
         // Act
@@ -506,24 +508,71 @@ await sutProvider.GetDependency<IEventService>()
             .Received(1)
             .LogUserEventAsync(user.Id, EventType.User_Disabled2fa);
 
-        // Revoke the user from the first organization because they are managed by it
+        // Revoke the user from the first organization
         await sutProvider.GetDependency<IRevokeNonCompliantOrganizationUserCommand>()
             .Received(1)
             .RevokeNonCompliantOrganizationUsersAsync(
                 Arg.Is<RevokeOrganizationUsersRequest>(r => r.OrganizationId == organization1.Id &&
-                    r.OrganizationUsers.First().UserId == user.Id &&
+                    r.OrganizationUsers.First().Id == organizationUserId1 &&
                     r.OrganizationUsers.First().OrganizationId == organization1.Id));
         await sutProvider.GetDependency<IMailService>()
             .Received(1)
             .SendOrganizationUserRevokedForTwoFactoryPolicyEmailAsync(organization1.DisplayName(), user.Email);
 
-        // Remove the user from the second organization because they are not managed by it
-        await sutProvider.GetDependency<IRemoveOrganizationUserCommand>()
+        // Remove the user from the second organization
+        await sutProvider.GetDependency<IRevokeNonCompliantOrganizationUserCommand>()
             .Received(1)
-            .RemoveUserAsync(organization2.Id, user.Id);
+            .RevokeNonCompliantOrganizationUsersAsync(
+                Arg.Is<RevokeOrganizationUsersRequest>(r => r.OrganizationId == organization2.Id &&
+                    r.OrganizationUsers.First().Id == organizationUserId2 &&
+                    r.OrganizationUsers.First().OrganizationId == organization2.Id));
         await sutProvider.GetDependency<IMailService>()
             .Received(1)
-            .SendOrganizationUserRemovedForPolicyTwoStepEmailAsync(organization2.DisplayName(), user.Email);
+            .SendOrganizationUserRevokedForTwoFactoryPolicyEmailAsync(organization2.DisplayName(), user.Email);
+    }
+
+    [Theory, BitAutoData]
+    public async Task DisableTwoFactorProviderAsync_WithAccountDeprovisioningEnabled_UserHasOneProviderEnabled_DoesNotRemoveUserFromOrganization(
+        SutProvider<UserService> sutProvider, User user, Organization organization)
+    {
+        // Arrange
+        user.SetTwoFactorProviders(new Dictionary<TwoFactorProviderType, TwoFactorProvider>
+        {
+            [TwoFactorProviderType.Email] = new() { Enabled = true },
+            [TwoFactorProviderType.Remember] = new() { Enabled = true }
+        });
+        sutProvider.GetDependency<IPolicyService>()
+            .GetPoliciesApplicableToUserAsync(user.Id, PolicyType.TwoFactorAuthentication)
+            .Returns(
+            [
+                new OrganizationUserPolicyDetails
+                {
+                    OrganizationId = organization.Id,
+                    PolicyType = PolicyType.TwoFactorAuthentication,
+                    PolicyEnabled = true
+                }
+            ]);
+        sutProvider.GetDependency<IOrganizationRepository>()
+            .GetByIdAsync(organization.Id)
+            .Returns(organization);
+        var expectedSavedProviders = JsonHelpers.LegacySerialize(new Dictionary<TwoFactorProviderType, TwoFactorProvider>
+        {
+            [TwoFactorProviderType.Remember] = new() { Enabled = true }
+        }, JsonHelpers.LegacyEnumKeyResolver);
+
+        // Act
+        await sutProvider.Sut.DisableTwoFactorProviderAsync(user, TwoFactorProviderType.Email);
+
+        // Assert
+        await sutProvider.GetDependency<IUserRepository>()
+            .Received(1)
+            .ReplaceAsync(Arg.Is<User>(u => u.Id == user.Id && u.TwoFactorProviders == expectedSavedProviders));
+        await sutProvider.GetDependency<IRevokeNonCompliantOrganizationUserCommand>()
+            .DidNotReceiveWithAnyArgs()
+            .RevokeNonCompliantOrganizationUsersAsync(default);
+        await sutProvider.GetDependency<IMailService>()
+            .DidNotReceiveWithAnyArgs()
+            .SendOrganizationUserRevokedForTwoFactoryPolicyEmailAsync(default, default);
     }
 
     [Theory, BitAutoData]