Added more tests to catch more use cases and fix bugs. (#5598)

src/Core/AdminConsole/OrganizationFeatures/OrganizationUsers/RestoreUser/v1/RestoreOrganizationUserCommand.cs

@@ -87,7 +87,10 @@ public class RestoreOrganizationUserCommand(
                 .twoFactorIsEnabled;
         }
 
-        await CheckUserForOtherFreeOrganizationOwnershipAsync(organizationUser);
+        if (organization.PlanType == PlanType.Free)
+        {
+            await CheckUserForOtherFreeOrganizationOwnershipAsync(organizationUser);
+        }
 
         await CheckPoliciesBeforeRestoreAsync(organizationUser, userTwoFactorIsEnabled);
 
@@ -100,7 +103,7 @@ public class RestoreOrganizationUserCommand(
 
     private async Task CheckUserForOtherFreeOrganizationOwnershipAsync(OrganizationUser organizationUser)
     {
-        var relatedOrgUsersFromOtherOrgs = await organizationUserRepository.GetManyByUserAsync(organizationUser.UserId.Value);
+        var relatedOrgUsersFromOtherOrgs = await organizationUserRepository.GetManyByUserAsync(organizationUser.UserId!.Value);
         var otherOrgs = await organizationRepository.GetManyByUserIdAsync(organizationUser.UserId.Value);
 
         var orgOrgUserDict = relatedOrgUsersFromOtherOrgs
@@ -110,13 +113,16 @@ public class RestoreOrganizationUserCommand(
         CheckForOtherFreeOrganizationOwnership(organizationUser, orgOrgUserDict);
     }
 
-    private async Task<Dictionary<OrganizationUser, Organization>> GetRelatedOrganizationUsersAndOrganizations(
+    private async Task<Dictionary<OrganizationUser, Organization>> GetRelatedOrganizationUsersAndOrganizationsAsync(
-        IEnumerable<OrganizationUser> organizationUsers)
+        List<OrganizationUser> organizationUsers)
     {
-        var allUserIds = organizationUsers.Select(x => x.UserId.Value);
+        var allUserIds = organizationUsers
+            .Where(x => x.UserId.HasValue)
+            .Select(x => x.UserId.Value);
 
         var otherOrganizationUsers = (await organizationUserRepository.GetManyByManyUsersAsync(allUserIds))
-            .Where(x => organizationUsers.Any(y => y.Id == x.Id) == false);
+            .Where(x => organizationUsers.Any(y => y.Id == x.Id) == false)
+            .ToArray();
 
         var otherOrgs = await organizationRepository.GetManyByIdsAsync(otherOrganizationUsers
                 .Select(x => x.OrganizationId)
@@ -130,7 +136,9 @@ public class RestoreOrganizationUserCommand(
         Dictionary<OrganizationUser, Organization> otherOrgUsersAndOrgs)
     {
         var ownerOrAdminList = new[] { OrganizationUserType.Owner, OrganizationUserType.Admin };
-        if (otherOrgUsersAndOrgs.Any(x =>
+
+        if (ownerOrAdminList.Any(x => organizationUser.Type == x) &&
+            otherOrgUsersAndOrgs.Any(x =>
                 x.Key.UserId == organizationUser.UserId &&
                 ownerOrAdminList.Any(userType => userType == x.Key.Type) &&
                 x.Key.Status == OrganizationUserStatusType.Confirmed &&
@@ -170,7 +178,7 @@ public class RestoreOrganizationUserCommand(
         var organizationUsersTwoFactorEnabled = await twoFactorIsEnabledQuery.TwoFactorIsEnabledAsync(
             filteredUsers.Where(ou => ou.UserId.HasValue).Select(ou => ou.UserId.Value));
 
-        var orgUsersAndOrgs = await GetRelatedOrganizationUsersAndOrganizations(filteredUsers);
+        var orgUsersAndOrgs = await GetRelatedOrganizationUsersAndOrganizationsAsync(filteredUsers);
 
         var result = new List<Tuple<OrganizationUser, string>>();
 
@@ -201,7 +209,10 @@ public class RestoreOrganizationUserCommand(
 
                 await CheckPoliciesBeforeRestoreAsync(organizationUser, twoFactorIsEnabled);
 
-                CheckForOtherFreeOrganizationOwnership(organizationUser, orgUsersAndOrgs);
+                if (organization.PlanType == PlanType.Free)
+                {
+                    CheckForOtherFreeOrganizationOwnership(organizationUser, orgUsersAndOrgs);
+                }
 
                 var status = OrganizationService.GetPriorActiveOrganizationUserStatusType(organizationUser);
 

test/Core.Test/AdminConsole/OrganizationFeatures/OrganizationUsers/RestoreUser/RestoreOrganizationUserCommandTests.cs

@@ -471,10 +471,11 @@ public class RestoreOrganizationUserCommandTests
         Organization organization,
         Organization otherOrganization,
         [OrganizationUser(OrganizationUserStatusType.Confirmed, OrganizationUserType.Owner)] OrganizationUser owner,
-        [OrganizationUser(OrganizationUserStatusType.Revoked)] OrganizationUser organizationUser,
+        [OrganizationUser(OrganizationUserStatusType.Revoked, OrganizationUserType.Owner)] OrganizationUser organizationUser,
         [OrganizationUser(OrganizationUserStatusType.Confirmed, OrganizationUserType.Owner)] OrganizationUser orgUserOwnerFromDifferentOrg,
         SutProvider<RestoreOrganizationUserCommand> sutProvider)
     {
+        organization.PlanType = PlanType.Free;
         organizationUser.Email = null; // this is required to mock that the user as had already been confirmed before the revoke
 
         orgUserOwnerFromDifferentOrg.UserId = organizationUser.UserId;
@@ -506,6 +507,107 @@ public class RestoreOrganizationUserCommandTests
         Assert.Equal("User is an owner/admin of another free organization. Please have them upgrade to a paid plan to restore their account.", exception.Message);
     }
 
+    [Theory, BitAutoData]
+    public async Task RestoreUser_WhenUserOwningAnotherFreeOrganizationAndIsOnlyAUserInCurrentOrg_ThenUserShouldBeRestored(
+        Organization organization,
+        Organization otherOrganization,
+        [OrganizationUser(OrganizationUserStatusType.Confirmed, OrganizationUserType.Owner)] OrganizationUser owner,
+        [OrganizationUser(OrganizationUserStatusType.Revoked)] OrganizationUser organizationUser,
+        [OrganizationUser(OrganizationUserStatusType.Confirmed, OrganizationUserType.Owner)] OrganizationUser orgUserOwnerFromDifferentOrg,
+        SutProvider<RestoreOrganizationUserCommand> sutProvider)
+    {
+        organization.PlanType = PlanType.Free;
+        organizationUser.Email = null; // this is required to mock that the user as had already been confirmed before the revoke
+
+        orgUserOwnerFromDifferentOrg.UserId = organizationUser.UserId;
+        otherOrganization.Id = orgUserOwnerFromDifferentOrg.OrganizationId;
+        otherOrganization.PlanType = PlanType.Free;
+
+        RestoreUser_Setup(organization, owner, organizationUser, sutProvider);
+
+        var organizationUserRepository = sutProvider.GetDependency<IOrganizationUserRepository>();
+        organizationUserRepository
+            .GetManyByUserAsync(organizationUser.UserId.Value)
+            .Returns([orgUserOwnerFromDifferentOrg]);
+
+        sutProvider.GetDependency<IOrganizationRepository>()
+            .GetManyByUserIdAsync(organizationUser.UserId.Value)
+            .Returns([otherOrganization]);
+
+        sutProvider.GetDependency<IPolicyService>()
+            .GetPoliciesApplicableToUserAsync(organizationUser.UserId.Value, PolicyType.TwoFactorAuthentication,
+                Arg.Any<OrganizationUserStatusType>())
+            .Returns([
+                new OrganizationUserPolicyDetails
+                {
+                    OrganizationId = organizationUser.OrganizationId,
+                    PolicyType = PolicyType.TwoFactorAuthentication
+                }
+            ]);
+
+        sutProvider.GetDependency<ITwoFactorIsEnabledQuery>()
+            .TwoFactorIsEnabledAsync(Arg.Is<IEnumerable<Guid>>(i => i.Contains(organizationUser.UserId.Value)))
+            .Returns(new List<(Guid userId, bool twoFactorIsEnabled)> { (organizationUser.UserId.Value, true) });
+
+        await sutProvider.Sut.RestoreUserAsync(organizationUser, owner.Id);
+
+        await organizationUserRepository
+            .Received(1)
+            .RestoreAsync(organizationUser.Id,
+                Arg.Is<OrganizationUserStatusType>(x => x != OrganizationUserStatusType.Revoked));
+    }
+
+    [Theory, BitAutoData]
+    public async Task RestoreUser_WhenUserOwningAnotherFreeOrganizationAndCurrentOrgIsNotFree_ThenUserShouldBeRestored(
+        Organization organization,
+        Organization otherOrganization,
+        [OrganizationUser(OrganizationUserStatusType.Confirmed, OrganizationUserType.Owner)] OrganizationUser owner,
+        [OrganizationUser(OrganizationUserStatusType.Revoked, OrganizationUserType.Owner)] OrganizationUser organizationUser,
+        [OrganizationUser(OrganizationUserStatusType.Confirmed, OrganizationUserType.Owner)] OrganizationUser orgUserOwnerFromDifferentOrg,
+        SutProvider<RestoreOrganizationUserCommand> sutProvider)
+    {
+        organization.PlanType = PlanType.EnterpriseAnnually2023;
+
+        organizationUser.Email = null; // this is required to mock that the user as had already been confirmed before the revoke
+
+        orgUserOwnerFromDifferentOrg.UserId = organizationUser.UserId;
+        otherOrganization.Id = orgUserOwnerFromDifferentOrg.OrganizationId;
+        otherOrganization.PlanType = PlanType.Free;
+
+        RestoreUser_Setup(organization, owner, organizationUser, sutProvider);
+
+        var organizationUserRepository = sutProvider.GetDependency<IOrganizationUserRepository>();
+        organizationUserRepository
+            .GetManyByUserAsync(organizationUser.UserId.Value)
+            .Returns([orgUserOwnerFromDifferentOrg]);
+
+        sutProvider.GetDependency<IOrganizationRepository>()
+            .GetManyByUserIdAsync(organizationUser.UserId.Value)
+            .Returns([otherOrganization]);
+
+        sutProvider.GetDependency<IPolicyService>()
+            .GetPoliciesApplicableToUserAsync(organizationUser.UserId.Value, PolicyType.TwoFactorAuthentication,
+                Arg.Any<OrganizationUserStatusType>())
+            .Returns([
+                new OrganizationUserPolicyDetails
+                {
+                    OrganizationId = organizationUser.OrganizationId,
+                    PolicyType = PolicyType.TwoFactorAuthentication
+                }
+            ]);
+
+        sutProvider.GetDependency<ITwoFactorIsEnabledQuery>()
+            .TwoFactorIsEnabledAsync(Arg.Is<IEnumerable<Guid>>(i => i.Contains(organizationUser.UserId.Value)))
+            .Returns(new List<(Guid userId, bool twoFactorIsEnabled)> { (organizationUser.UserId.Value, true) });
+
+        await sutProvider.Sut.RestoreUserAsync(organizationUser, owner.Id);
+
+        await organizationUserRepository
+            .Received(1)
+            .RestoreAsync(organizationUser.Id,
+                Arg.Is<OrganizationUserStatusType>(x => x != OrganizationUserStatusType.Revoked));
+    }
+
     [Theory, BitAutoData]
     public async Task RestoreUsers_Success(Organization organization,
     [OrganizationUser(OrganizationUserStatusType.Confirmed, OrganizationUserType.Owner)] OrganizationUser owner,
@@ -612,7 +714,7 @@ public class RestoreOrganizationUserCommandTests
     [Theory, BitAutoData]
     public async Task RestoreUsers_UserOwnsAnotherFreeOrganization_BlocksOwnerUserFromBeingRestored(Organization organization,
         [OrganizationUser(OrganizationUserStatusType.Confirmed, OrganizationUserType.Owner)] OrganizationUser owner,
-        [OrganizationUser(OrganizationUserStatusType.Revoked)] OrganizationUser orgUser1,
+        [OrganizationUser(OrganizationUserStatusType.Revoked, OrganizationUserType.Owner)] OrganizationUser orgUser1,
         [OrganizationUser(OrganizationUserStatusType.Revoked)] OrganizationUser orgUser2,
         [OrganizationUser(OrganizationUserStatusType.Revoked)] OrganizationUser orgUser3,
         [OrganizationUser(OrganizationUserStatusType.Confirmed, OrganizationUserType.Owner)] OrganizationUser orgUserFromOtherOrg,
@@ -637,7 +739,7 @@ public class RestoreOrganizationUserCommandTests
 
         organizationUserRepository
             .GetManyAsync(Arg.Is<IEnumerable<Guid>>(ids => ids.Contains(orgUser1.Id) && ids.Contains(orgUser2.Id) && ids.Contains(orgUser3.Id)))
-            .Returns(new[] { orgUser1, orgUser2, orgUser3 });
+            .Returns([orgUser1, orgUser2, orgUser3]);
 
         userRepository.GetByIdAsync(orgUser2.UserId!.Value).Returns(new User { Email = "test@example.com" });
 
@@ -674,6 +776,110 @@ public class RestoreOrganizationUserCommandTests
             .RestoreAsync(orgUser1.Id, OrganizationUserStatusType.Confirmed);
     }
 
+    [Theory, BitAutoData]
+    public async Task RestoreUsers_UserOwnsAnotherFreeOrganizationButReactivatingOrgIsPaid_RestoresUser(Organization organization,
+        [OrganizationUser(OrganizationUserStatusType.Confirmed, OrganizationUserType.Owner)] OrganizationUser owner,
+        [OrganizationUser(OrganizationUserStatusType.Revoked, OrganizationUserType.Owner)] OrganizationUser orgUser1,
+        [OrganizationUser(OrganizationUserStatusType.Confirmed, OrganizationUserType.Owner)] OrganizationUser orgUserFromOtherOrg,
+        Organization otherOrganization,
+        SutProvider<RestoreOrganizationUserCommand> sutProvider)
+    {
+        // Arrange
+        organization.PlanType = PlanType.EnterpriseAnnually2023;
+
+        RestoreUser_Setup(organization, owner, orgUser1, sutProvider);
+        var organizationUserRepository = sutProvider.GetDependency<IOrganizationUserRepository>();
+        var policyService = sutProvider.GetDependency<IPolicyService>();
+        var userService = Substitute.For<IUserService>();
+
+        orgUser1.OrganizationId = organization.Id;
+
+        orgUserFromOtherOrg.UserId = orgUser1.UserId;
+
+        otherOrganization.Id = orgUserFromOtherOrg.OrganizationId;
+        otherOrganization.PlanType = PlanType.Free;
+
+        organizationUserRepository
+            .GetManyAsync(Arg.Is<IEnumerable<Guid>>(ids => ids.Contains(orgUser1.Id)))
+            .Returns([orgUser1]);
+
+        organizationUserRepository
+            .GetManyByManyUsersAsync(Arg.Any<IEnumerable<Guid>>())
+            .Returns([orgUserFromOtherOrg]);
+
+        sutProvider.GetDependency<IOrganizationRepository>()
+            .GetManyByIdsAsync(Arg.Is<IEnumerable<Guid>>(ids => ids.Contains(orgUserFromOtherOrg.OrganizationId)))
+            .Returns([otherOrganization]);
+
+
+        // Setup 2FA policy
+        policyService.GetPoliciesApplicableToUserAsync(Arg.Any<Guid>(), PolicyType.TwoFactorAuthentication, Arg.Any<OrganizationUserStatusType>())
+            .Returns([new OrganizationUserPolicyDetails { OrganizationId = organization.Id, PolicyType = PolicyType.TwoFactorAuthentication }]);
+
+        // User1 has 2FA, User2 doesn't
+        sutProvider.GetDependency<ITwoFactorIsEnabledQuery>()
+            .TwoFactorIsEnabledAsync(Arg.Is<IEnumerable<Guid>>(ids => ids.Contains(orgUser1.UserId!.Value)))
+            .Returns(new List<(Guid userId, bool twoFactorIsEnabled)>
+            {
+                (orgUser1.UserId!.Value, true)
+            });
+
+        // Act
+        var result = await sutProvider.Sut.RestoreUsersAsync(organization.Id, [orgUser1.Id], owner.Id, userService);
+
+        // Assert
+        Assert.Single(result);
+        Assert.Equal(string.Empty, result[0].Item2);
+        await organizationUserRepository
+            .Received(1)
+            .RestoreAsync(orgUser1.Id, Arg.Is<OrganizationUserStatusType>(x => x != OrganizationUserStatusType.Revoked));
+    }
+
+    [Theory]
+    [BitAutoData]
+    public async Task RestoreUsers_UserOwnsAnotherOrganizationButIsOnlyUserOfCurrentOrganization_UserShouldBeRestored(
+        Organization organization,
+        [OrganizationUser(OrganizationUserStatusType.Confirmed, OrganizationUserType.Owner)] OrganizationUser owner,
+        [OrganizationUser(OrganizationUserStatusType.Revoked, OrganizationUserType.User)] OrganizationUser orgUser1,
+        [OrganizationUser(OrganizationUserStatusType.Confirmed, OrganizationUserType.Owner)] OrganizationUser orgUserFromOtherOrg,
+        Organization otherOrganization,
+        SutProvider<RestoreOrganizationUserCommand> sutProvider)
+    {
+        // Arrange
+        organization.PlanType = PlanType.Free;
+
+        RestoreUser_Setup(organization, owner, orgUser1, sutProvider);
+        var organizationUserRepository = sutProvider.GetDependency<IOrganizationUserRepository>();
+        var userService = Substitute.For<IUserService>();
+
+        orgUser1.OrganizationId = organization.Id;
+
+        orgUserFromOtherOrg.UserId = orgUser1.UserId;
+
+        otherOrganization.Id = orgUserFromOtherOrg.OrganizationId;
+        otherOrganization.PlanType = PlanType.Free;
+
+        organizationUserRepository
+            .GetManyAsync(Arg.Is<IEnumerable<Guid>>(ids => ids.Contains(orgUser1.Id)))
+            .Returns([orgUser1]);
+
+        organizationUserRepository
+            .GetManyByManyUsersAsync(Arg.Any<IEnumerable<Guid>>())
+            .Returns([orgUserFromOtherOrg]);
+
+        sutProvider.GetDependency<IPolicyService>().GetPoliciesApplicableToUserAsync(Arg.Any<Guid>(), PolicyType.TwoFactorAuthentication, Arg.Any<OrganizationUserStatusType>())
+            .Returns([new OrganizationUserPolicyDetails { OrganizationId = organization.Id, PolicyType = PolicyType.TwoFactorAuthentication }]);
+
+        // Act
+        var result = await sutProvider.Sut.RestoreUsersAsync(organization.Id, [orgUser1.Id], owner.Id, userService);
+
+        Assert.Single(result);
+        Assert.Equal(string.Empty, result[0].Item2);
+        await organizationUserRepository
+            .Received(1)
+            .RestoreAsync(orgUser1.Id, Arg.Is<OrganizationUserStatusType>(x => x != OrganizationUserStatusType.Revoked));
+    }
+
     private static void RestoreUser_Setup(
         Organization organization,
         OrganizationUser? requestingOrganizationUser,