[PM-29246] Simplify Usage of Organization Policies (#6837)

* Initial implementation of new policy query * Remove unused using * Adjusts method name to better match repository method * Correct namespace * Initial refactor of policy loading * Add xml doc, incorporate shim data model * Updates usages to reflect new shim model * Prune extranneous data from policy detail response model, format code * Fix broken test, delete inapplicable test * Adds test cases covering query * Adjust codebase to use new PolicyQueryçˆ * Format code * Fix incorrect mock on test * Fix formatting * Adjust method name * More naming adjustments * Add PolicyData constructor, update test usages * Rename PolicyData -> PolicyStatus * Remove unused using
2026-01-31 14:13:18 +08:00 · 2026-01-29 14:11:20 -06:00
parent 0544ec41d5
commit 93a28eed40
33 changed files with 457 additions and 302 deletions
--- a/test/Core.Test/OrganizationFeatures/OrganizationSubscriptionUpdate/UpgradeOrganizationPlanCommandTests.cs
+++ b/test/Core.Test/OrganizationFeatures/OrganizationSubscriptionUpdate/UpgradeOrganizationPlanCommandTests.cs
@@ -1,4 +1,7 @@
-using Bit.Core.Billing.Enums;
+using Bit.Core.AdminConsole.Enums;
+using Bit.Core.AdminConsole.Models.Data.Organizations.Policies;
+using Bit.Core.AdminConsole.OrganizationFeatures.Policies;
+using Bit.Core.Billing.Enums;
 using Bit.Core.Billing.Pricing;
 using Bit.Core.Billing.Services;
 using Bit.Core.Exceptions;
@@ -9,6 +12,7 @@ using Bit.Core.OrganizationFeatures.OrganizationSubscriptions;
 using Bit.Core.Repositories;
 using Bit.Core.SecretsManager.Repositories;
 using Bit.Core.Services;
+using Bit.Core.Test.AdminConsole.AutoFixture;
 using Bit.Core.Test.AutoFixture.OrganizationFixtures;
 using Bit.Core.Test.Billing.Mocks;
 using Bit.Test.Common.AutoFixture;
@@ -72,8 +76,12 @@ public class UpgradeOrganizationPlanCommandTests
    [Theory]
    [FreeOrganizationUpgradeCustomize, BitAutoData]
    public async Task UpgradePlan_Passes(Organization organization, OrganizationUpgrade upgrade,
+        [Policy(PolicyType.ResetPassword, false)] PolicyStatus policy,
            SutProvider<UpgradeOrganizationPlanCommand> sutProvider)
    {
+        sutProvider.GetDependency<IPolicyQuery>()
+            .RunAsync(Arg.Any<Guid>(), Arg.Any<PolicyType>())
+            .Returns(policy);
        sutProvider.GetDependency<IOrganizationRepository>().GetByIdAsync(organization.Id).Returns(organization);
        sutProvider.GetDependency<IPricingClient>().GetPlanOrThrow(organization.PlanType).Returns(MockPlans.Get(organization.PlanType));
        upgrade.AdditionalSmSeats = 10;
@@ -100,6 +108,7 @@ public class UpgradeOrganizationPlanCommandTests
        PlanType planType,
        Organization organization,
        OrganizationUpgrade organizationUpgrade,
+        [Policy(PolicyType.ResetPassword, false)] PolicyStatus policy,
        SutProvider<UpgradeOrganizationPlanCommand> sutProvider)
    {
        sutProvider.GetDependency<IOrganizationRepository>().GetByIdAsync(organization.Id).Returns(organization);
@@ -116,6 +125,9 @@ public class UpgradeOrganizationPlanCommandTests
        organizationUpgrade.Plan = planType;

        sutProvider.GetDependency<IPricingClient>().GetPlanOrThrow(organizationUpgrade.Plan).Returns(MockPlans.Get(organizationUpgrade.Plan));
+        sutProvider.GetDependency<IPolicyQuery>()
+            .RunAsync(Arg.Any<Guid>(), Arg.Any<PolicyType>())
+            .Returns(policy);
        sutProvider.GetDependency<IOrganizationRepository>()
            .GetOccupiedSeatCountByOrganizationIdAsync(organization.Id).Returns(new OrganizationSeatCounts
            {
@@ -141,15 +153,20 @@ public class UpgradeOrganizationPlanCommandTests
    [BitAutoData(PlanType.TeamsAnnually)]
    [BitAutoData(PlanType.TeamsStarter)]
    public async Task UpgradePlan_SM_Passes(PlanType planType, Organization organization, OrganizationUpgrade upgrade,
+        [Policy(PolicyType.ResetPassword, false)] PolicyStatus policy,
        SutProvider<UpgradeOrganizationPlanCommand> sutProvider)
    {
-        sutProvider.GetDependency<IPricingClient>().GetPlanOrThrow(organization.PlanType).Returns(MockPlans.Get(organization.PlanType));

        upgrade.Plan = planType;
        sutProvider.GetDependency<IPricingClient>().GetPlanOrThrow(upgrade.Plan).Returns(MockPlans.Get(upgrade.Plan));

        var plan = MockPlans.Get(upgrade.Plan);

+        sutProvider.GetDependency<IPolicyQuery>()
+            .RunAsync(Arg.Any<Guid>(), Arg.Any<PolicyType>())
+            .Returns(policy);
+
+        sutProvider.GetDependency<IPricingClient>().GetPlanOrThrow(organization.PlanType).Returns(MockPlans.Get(organization.PlanType));
        sutProvider.GetDependency<IOrganizationRepository>().GetByIdAsync(organization.Id).Returns(organization);

        upgrade.AdditionalSeats = 15;
@@ -180,6 +197,7 @@ public class UpgradeOrganizationPlanCommandTests
    [BitAutoData(PlanType.TeamsAnnually)]
    [BitAutoData(PlanType.TeamsStarter)]
    public async Task UpgradePlan_SM_NotEnoughSmSeats_Throws(PlanType planType, Organization organization, OrganizationUpgrade upgrade,
+        [Policy(PolicyType.ResetPassword, false)] PolicyStatus policy,
        SutProvider<UpgradeOrganizationPlanCommand> sutProvider)
    {
        upgrade.Plan = planType;
@@ -191,6 +209,10 @@ public class UpgradeOrganizationPlanCommandTests
        organization.SmSeats = 2;
        sutProvider.GetDependency<IPricingClient>().GetPlanOrThrow(organization.PlanType).Returns(MockPlans.Get(organization.PlanType));

+        sutProvider.GetDependency<IPolicyQuery>()
+            .RunAsync(Arg.Any<Guid>(), Arg.Any<PolicyType>())
+            .Returns(policy);
+
        sutProvider.GetDependency<IOrganizationRepository>().GetByIdAsync(organization.Id).Returns(organization);
        sutProvider.GetDependency<IOrganizationRepository>()
            .GetOccupiedSeatCountByOrganizationIdAsync(organization.Id).Returns(new OrganizationSeatCounts
@@ -214,7 +236,9 @@ public class UpgradeOrganizationPlanCommandTests
    [BitAutoData(PlanType.TeamsAnnually, 51)]
    [BitAutoData(PlanType.TeamsStarter, 51)]
    public async Task UpgradePlan_SM_NotEnoughServiceAccounts_Throws(PlanType planType, int currentServiceAccounts,
-     Organization organization, OrganizationUpgrade upgrade, SutProvider<UpgradeOrganizationPlanCommand> sutProvider)
+     Organization organization, OrganizationUpgrade upgrade,
+     [Policy(PolicyType.ResetPassword, false)] PolicyStatus policy,
+     SutProvider<UpgradeOrganizationPlanCommand> sutProvider)
    {
        upgrade.Plan = planType;
        upgrade.AdditionalSeats = 15;
@@ -226,6 +250,10 @@ public class UpgradeOrganizationPlanCommandTests
        organization.SmServiceAccounts = currentServiceAccounts;
        sutProvider.GetDependency<IPricingClient>().GetPlanOrThrow(organization.PlanType).Returns(MockPlans.Get(organization.PlanType));

+        sutProvider.GetDependency<IPolicyQuery>()
+            .RunAsync(Arg.Any<Guid>(), Arg.Any<PolicyType>())
+            .Returns(policy);
+
        sutProvider.GetDependency<IOrganizationRepository>().GetByIdAsync(organization.Id).Returns(organization);
        sutProvider.GetDependency<IOrganizationRepository>()
            .GetOccupiedSeatCountByOrganizationIdAsync(organization.Id).Returns(new OrganizationSeatCounts
@@ -251,6 +279,7 @@ public class UpgradeOrganizationPlanCommandTests
        OrganizationUpgrade upgrade,
        string newPublicKey,
        string newPrivateKey,
+        [Policy(PolicyType.ResetPassword, false)] PolicyStatus policy,
        SutProvider<UpgradeOrganizationPlanCommand> sutProvider)
    {
        organization.PublicKey = null;
@@ -262,6 +291,9 @@ public class UpgradeOrganizationPlanCommandTests
            publicKey: newPublicKey);
        upgrade.AdditionalSeats = 10;

+        sutProvider.GetDependency<IPolicyQuery>()
+            .RunAsync(Arg.Any<Guid>(), Arg.Any<PolicyType>())
+            .Returns(policy);
        sutProvider.GetDependency<IOrganizationRepository>()
            .GetByIdAsync(organization.Id)
            .Returns(organization);
@@ -291,6 +323,7 @@ public class UpgradeOrganizationPlanCommandTests
    public async Task UpgradePlan_WhenOrganizationAlreadyHasPublicAndPrivateKeys_DoesNotOverwriteWithNull(
        Organization organization,
        OrganizationUpgrade upgrade,
+        [Policy(PolicyType.ResetPassword, false)] PolicyStatus policy,
        SutProvider<UpgradeOrganizationPlanCommand> sutProvider)
    {
        // Arrange
@@ -304,6 +337,9 @@ public class UpgradeOrganizationPlanCommandTests
        upgrade.Keys = null;
        upgrade.AdditionalSeats = 10;

+        sutProvider.GetDependency<IPolicyQuery>()
+            .RunAsync(Arg.Any<Guid>(), Arg.Any<PolicyType>())
+            .Returns(policy);
        sutProvider.GetDependency<IOrganizationRepository>()
            .GetByIdAsync(organization.Id)
            .Returns(organization);
@@ -333,6 +369,7 @@ public class UpgradeOrganizationPlanCommandTests
    public async Task UpgradePlan_WhenOrganizationAlreadyHasPublicAndPrivateKeys_DoesNotBackfillWithNewKeys(
        Organization organization,
        OrganizationUpgrade upgrade,
+        [Policy(PolicyType.ResetPassword, false)] PolicyStatus policy,
        SutProvider<UpgradeOrganizationPlanCommand> sutProvider)
    {
        // Arrange
@@ -343,6 +380,9 @@ public class UpgradeOrganizationPlanCommandTests

        organization.PublicKey = existingPublicKey;
        organization.PrivateKey = existingPrivateKey;
+        sutProvider.GetDependency<IPolicyQuery>()
+            .RunAsync(Arg.Any<Guid>(), Arg.Any<PolicyType>())
+            .Returns(policy);

        upgrade.Plan = PlanType.TeamsAnnually;
        upgrade.Keys = new PublicKeyEncryptionKeyPairData(
--- a/test/Core.Test/OrganizationFeatures/Policies/PolicyQueryTests.cs
+++ b/test/Core.Test/OrganizationFeatures/Policies/PolicyQueryTests.cs
@@ -0,0 +1,55 @@
+using Bit.Core.AdminConsole.Entities;
+using Bit.Core.AdminConsole.Enums;
+using Bit.Core.AdminConsole.OrganizationFeatures.Policies.Implementations;
+using Bit.Core.AdminConsole.Repositories;
+using Bit.Test.Common.AutoFixture;
+using Bit.Test.Common.AutoFixture.Attributes;
+using NSubstitute;
+using NSubstitute.ReturnsExtensions;
+using Xunit;
+
+namespace Bit.Core.Test.OrganizationFeatures.Policies;
+
+[SutProviderCustomize]
+public class PolicyQueryTests
+{
+    [Theory, BitAutoData]
+    public async Task RunAsync_WithExistingPolicy_ReturnsPolicy(SutProvider<PolicyQuery> sutProvider,
+        Policy policy)
+    {
+        // Arrange
+        sutProvider.GetDependency<IPolicyRepository>()
+            .GetByOrganizationIdTypeAsync(policy.OrganizationId, policy.Type)
+            .Returns(policy);
+
+        // Act
+        var policyData = await sutProvider.Sut.RunAsync(policy.OrganizationId, policy.Type);
+
+        // Assert
+        Assert.Equal(policy.Data, policyData.Data);
+        Assert.Equal(policy.Type, policyData.Type);
+        Assert.Equal(policy.Enabled, policyData.Enabled);
+        Assert.Equal(policy.OrganizationId, policyData.OrganizationId);
+    }
+
+    [Theory, BitAutoData]
+    public async Task RunAsync_WithNonExistentPolicy_ReturnsDefaultDisabledPolicy(
+        SutProvider<PolicyQuery> sutProvider,
+        Guid organizationId,
+        PolicyType policyType)
+    {
+        // Arrange
+        sutProvider.GetDependency<IPolicyRepository>()
+            .GetByOrganizationIdTypeAsync(organizationId, policyType)
+            .ReturnsNull();
+
+        // Act
+        var policyData = await sutProvider.Sut.RunAsync(organizationId, policyType);
+
+        // Assert
+        Assert.Equal(organizationId, policyData.OrganizationId);
+        Assert.Equal(policyType, policyData.Type);
+        Assert.False(policyData.Enabled);
+        Assert.Null(policyData.Data);
+    }
+}