From 23d9e3f4e581874353c65bc587cb87806312ae51 Mon Sep 17 00:00:00 2001
From: Thomas Avery <tavery@bitwarden.com>
Date: Fri, 25 Oct 2024 11:03:12 -0500
Subject: [PATCH] update tests

---
 .../AccountsKeyManagementControllerTests.cs   | 43 +++++++++++++------
 .../AccountsKeyManagementControllerTests.cs   | 17 +++++---
 ...egenerateUserAsymmetricKeysCommandTests.cs |  1 +
 3 files changed, 43 insertions(+), 18 deletions(-)

diff --git a/test/Api.IntegrationTest/KeyManagement/Controllers/AccountsKeyManagementControllerTests.cs b/test/Api.IntegrationTest/KeyManagement/Controllers/AccountsKeyManagementControllerTests.cs
index 0895bb1d9..ec7ca3746 100644
--- a/test/Api.IntegrationTest/KeyManagement/Controllers/AccountsKeyManagementControllerTests.cs
+++ b/test/Api.IntegrationTest/KeyManagement/Controllers/AccountsKeyManagementControllerTests.cs
@@ -19,6 +19,7 @@ public class AccountsKeyManagementControllerTests : IClassFixture<ApiApplication
 
     private readonly HttpClient _client;
     private readonly IEmergencyAccessRepository _emergencyAccessRepository;
+    private readonly IOrganizationUserRepository _organizationUserRepository;
     private readonly ApiApplicationFactory _factory;
     private readonly LoginHelper _loginHelper;
     private readonly IUserRepository _userRepository;
@@ -33,6 +34,7 @@ public class AccountsKeyManagementControllerTests : IClassFixture<ApiApplication
         _loginHelper = new LoginHelper(_factory, _client);
         _userRepository = _factory.GetService<IUserRepository>();
         _emergencyAccessRepository = _factory.GetService<IEmergencyAccessRepository>();
+        _organizationUserRepository = _factory.GetService<IOrganizationUserRepository>();
     }
 
     public async Task InitializeAsync()
@@ -80,24 +82,30 @@ public class AccountsKeyManagementControllerTests : IClassFixture<ApiApplication
     }
 
     [Theory]
-    [BitAutoData(true, true)]
-    [BitAutoData(true, false)]
-    [BitAutoData(false, true)]
+    [BitAutoData(OrganizationUserStatusType.Confirmed, EmergencyAccessStatusType.Confirmed)]
+    [BitAutoData(OrganizationUserStatusType.Confirmed, EmergencyAccessStatusType.RecoveryApproved)]
+    [BitAutoData(OrganizationUserStatusType.Confirmed, EmergencyAccessStatusType.RecoveryInitiated)]
+    [BitAutoData(OrganizationUserStatusType.Revoked, EmergencyAccessStatusType.Confirmed)]
+    [BitAutoData(OrganizationUserStatusType.Revoked, EmergencyAccessStatusType.RecoveryApproved)]
+    [BitAutoData(OrganizationUserStatusType.Revoked, EmergencyAccessStatusType.RecoveryInitiated)]
+    [BitAutoData(OrganizationUserStatusType.Confirmed, null)]
+    [BitAutoData(OrganizationUserStatusType.Revoked, null)]
+    [BitAutoData(OrganizationUserStatusType.Invited, EmergencyAccessStatusType.Confirmed)]
+    [BitAutoData(OrganizationUserStatusType.Invited, EmergencyAccessStatusType.RecoveryApproved)]
+    [BitAutoData(OrganizationUserStatusType.Invited, EmergencyAccessStatusType.RecoveryInitiated)]
     public async Task RegenerateKeysAsync_UserInOrgOrHasDesignatedEmergencyAccess_ThrowsBadRequest(
-        bool inOrganization,
-        bool hasDesignatedEmergencyAccess,
+        OrganizationUserStatusType organizationUserStatus,
+        EmergencyAccessStatusType? emergencyAccessStatus,
         KeyRegenerationRequestModel request)
     {
-        if (inOrganization)
+        if (organizationUserStatus is OrganizationUserStatusType.Confirmed or OrganizationUserStatusType.Revoked)
         {
-            await OrganizationTestHelpers.SignUpAsync(_factory,
-                PlanType.EnterpriseAnnually, _ownerEmail, passwordManagerSeats: 10,
-                paymentMethod: PaymentMethodType.Card);
+            await CreateOrganizationUserAsync(organizationUserStatus);
         }
 
-        if (hasDesignatedEmergencyAccess)
+        if (emergencyAccessStatus != null)
         {
-            await CreateDesignatedEmergencyAccessAsync();
+            await CreateDesignatedEmergencyAccessAsync(emergencyAccessStatus.Value);
         }
 
         await _loginHelper.LoginAsync(_ownerEmail);
@@ -124,7 +132,16 @@ public class AccountsKeyManagementControllerTests : IClassFixture<ApiApplication
         Assert.Equal(request.UserKeyEncryptedUserPrivateKey, user.PrivateKey);
     }
 
-    private async Task CreateDesignatedEmergencyAccessAsync()
+    private async Task CreateOrganizationUserAsync(OrganizationUserStatusType organizationUserStatus)
+    {
+        var (_, organizationUser) = await OrganizationTestHelpers.SignUpAsync(_factory,
+            PlanType.EnterpriseAnnually, _ownerEmail, passwordManagerSeats: 10,
+            paymentMethod: PaymentMethodType.Card);
+        organizationUser.Status = organizationUserStatus;
+        await _organizationUserRepository.ReplaceAsync(organizationUser);
+    }
+
+    private async Task CreateDesignatedEmergencyAccessAsync(EmergencyAccessStatusType emergencyAccessStatus)
     {
         var tempEmail = $"integration-test{Guid.NewGuid()}@bitwarden.com";
         await _factory.LoginWithNewAccount(tempEmail);
@@ -136,7 +153,7 @@ public class AccountsKeyManagementControllerTests : IClassFixture<ApiApplication
             GrantorId = tempUser!.Id,
             GranteeId = user!.Id,
             KeyEncrypted = _mockEncryptedString,
-            Status = EmergencyAccessStatusType.Confirmed,
+            Status = emergencyAccessStatus,
             Type = EmergencyAccessType.View,
             WaitTimeDays = 10,
             CreationDate = DateTime.UtcNow,
diff --git a/test/Api.Test/KeyManagement/Controllers/AccountsKeyManagementControllerTests.cs b/test/Api.Test/KeyManagement/Controllers/AccountsKeyManagementControllerTests.cs
index d01215b32..2615697ad 100644
--- a/test/Api.Test/KeyManagement/Controllers/AccountsKeyManagementControllerTests.cs
+++ b/test/Api.Test/KeyManagement/Controllers/AccountsKeyManagementControllerTests.cs
@@ -69,11 +69,15 @@ public class AccountsKeyManagementControllerTests
     [Theory]
     [BitAutoData]
     public async Task RegenerateKeysAsync_Success(SutProvider<AccountsKeyManagementController> sutProvider,
-        KeyRegenerationRequestModel data, User user)
+        KeyRegenerationRequestModel data, User user, ICollection<OrganizationUser> orgUsers,
+        ICollection<EmergencyAccessDetails> accessDetails)
     {
         sutProvider.GetDependency<IFeatureService>().IsEnabled(Arg.Is(FeatureFlagKeys.PrivateKeyRegeneration))
             .Returns(true);
         sutProvider.GetDependency<IUserService>().GetUserByPrincipalAsync(Arg.Any<ClaimsPrincipal>()).Returns(user);
+        sutProvider.GetDependency<IOrganizationUserRepository>().GetManyByUserAsync(Arg.Is(user.Id)).Returns(orgUsers);
+        sutProvider.GetDependency<IEmergencyAccessRepository>().GetManyDetailsByGranteeIdAsync(Arg.Is(user.Id))
+            .Returns(accessDetails);
 
         await sutProvider.Sut.RegenerateKeysAsync(data);
 
@@ -81,9 +85,12 @@ public class AccountsKeyManagementControllerTests
             .GetManyByUserAsync(Arg.Is(user.Id));
         await sutProvider.GetDependency<IEmergencyAccessRepository>().Received(1)
             .GetManyDetailsByGranteeIdAsync(Arg.Is(user.Id));
-        await sutProvider.GetDependency<IRegenerateUserAsymmetricKeysCommand>().ReceivedWithAnyArgs(1)
-            .RegenerateKeysAsync(Arg.Any<UserAsymmetricKeys>(),
-                Arg.Any<ICollection<OrganizationUser>>(),
-                Arg.Any<ICollection<EmergencyAccessDetails>>());
+        await sutProvider.GetDependency<IRegenerateUserAsymmetricKeysCommand>().Received(1)
+            .RegenerateKeysAsync(
+                Arg.Is<UserAsymmetricKeys>(u =>
+                    u.UserId == user.Id && u.PublicKey == data.UserPublicKey &&
+                    u.UserKeyEncryptedPrivateKey == data.UserKeyEncryptedUserPrivateKey),
+                Arg.Is(orgUsers),
+                Arg.Is(accessDetails));
     }
 }
diff --git a/test/Core.Test/KeyManagement/Commands/RegenerateUserAsymmetricKeysCommandTests.cs b/test/Core.Test/KeyManagement/Commands/RegenerateUserAsymmetricKeysCommandTests.cs
index f35673e0e..ce2b9afef 100644
--- a/test/Core.Test/KeyManagement/Commands/RegenerateUserAsymmetricKeysCommandTests.cs
+++ b/test/Core.Test/KeyManagement/Commands/RegenerateUserAsymmetricKeysCommandTests.cs
@@ -56,6 +56,7 @@ public class RegenerateUserAsymmetricKeysCommandTests
     [BitAutoData(false, true, false)]
     [BitAutoData(false, true, true)]
     [BitAutoData(true, false, false)]
+    [BitAutoData(true, false, true)]
     [BitAutoData(true, true, false)]
     [BitAutoData(true, true, true)]
     public async Task RegenerateKeysAsync_UserIdMisMatch_NotFoundException(