From 077d0fa6d7268a402dcbe88e574d2c4242cd6c78 Mon Sep 17 00:00:00 2001 From: Conner Turnbull <133619638+cturnbull-bitwarden@users.noreply.github.com> Date: Fri, 2 May 2025 12:53:06 -0400 Subject: [PATCH] Resolved an issue where autoscaling always happened (#5765) --- .../Services/IOrganizationService.cs | 1 + .../Implementations/OrganizationService.cs | 2 +- .../CreateSponsorshipCommand.cs | 19 ++- .../CreateSponsorshipCommandTests.cs | 127 ++++++++++++++++++ 4 files changed, 145 insertions(+), 4 deletions(-) diff --git a/src/Core/AdminConsole/Services/IOrganizationService.cs b/src/Core/AdminConsole/Services/IOrganizationService.cs index 9c9e311a02..1e53be734e 100644 --- a/src/Core/AdminConsole/Services/IOrganizationService.cs +++ b/src/Core/AdminConsole/Services/IOrganizationService.cs @@ -49,6 +49,7 @@ public interface IOrganizationService IEnumerable organizationUserIds, Guid? revokingUserId); Task CreatePendingOrganization(Organization organization, string ownerEmail, ClaimsPrincipal user, IUserService userService, bool salesAssistedTrialStarted); Task ReplaceAndUpdateCacheAsync(Organization org, EventType? orgEvent = null); + Task<(bool canScale, string failureReason)> CanScaleAsync(Organization organization, int seatsToAdd); void ValidatePasswordManagerPlan(Models.StaticStore.Plan plan, OrganizationUpgrade upgrade); void ValidateSecretsManagerPlan(Models.StaticStore.Plan plan, OrganizationUpgrade upgrade); diff --git a/src/Core/AdminConsole/Services/Implementations/OrganizationService.cs b/src/Core/AdminConsole/Services/Implementations/OrganizationService.cs index 532aebf5e0..5c7e5e29ed 100644 --- a/src/Core/AdminConsole/Services/Implementations/OrganizationService.cs +++ b/src/Core/AdminConsole/Services/Implementations/OrganizationService.cs @@ -1058,7 +1058,7 @@ public class OrganizationService : IOrganizationService organization: organization, initOrganization: initOrganization)); - internal async Task<(bool canScale, string failureReason)> CanScaleAsync( + public async Task<(bool canScale, string failureReason)> CanScaleAsync( Organization organization, int seatsToAdd) { diff --git a/src/Core/OrganizationFeatures/OrganizationSponsorships/FamiliesForEnterprise/CreateSponsorshipCommand.cs b/src/Core/OrganizationFeatures/OrganizationSponsorships/FamiliesForEnterprise/CreateSponsorshipCommand.cs index 3b74baf6f9..b15cbea240 100644 --- a/src/Core/OrganizationFeatures/OrganizationSponsorships/FamiliesForEnterprise/CreateSponsorshipCommand.cs +++ b/src/Core/OrganizationFeatures/OrganizationSponsorships/FamiliesForEnterprise/CreateSponsorshipCommand.cs @@ -15,7 +15,8 @@ public class CreateSponsorshipCommand( ICurrentContext currentContext, IOrganizationSponsorshipRepository organizationSponsorshipRepository, IUserService userService, - IOrganizationService organizationService) : ICreateSponsorshipCommand + IOrganizationService organizationService, + IOrganizationUserRepository organizationUserRepository) : ICreateSponsorshipCommand { public async Task CreateSponsorshipAsync( Organization sponsoringOrganization, @@ -82,14 +83,26 @@ public class CreateSponsorshipCommand( if (existingOrgSponsorship != null) { - // Replace existing invalid offer with our new sponsorship offer sponsorship.Id = existingOrgSponsorship.Id; } } if (isAdminInitiated && sponsoringOrganization.Seats.HasValue) { - await organizationService.AutoAddSeatsAsync(sponsoringOrganization, 1); + var occupiedSeats = await organizationUserRepository.GetOccupiedSeatCountByOrganizationIdAsync(sponsoringOrganization.Id); + var availableSeats = sponsoringOrganization.Seats.Value - occupiedSeats; + + if (availableSeats <= 0) + { + var newSeatsRequired = 1; + var (canScale, failureReason) = await organizationService.CanScaleAsync(sponsoringOrganization, newSeatsRequired); + if (!canScale) + { + throw new BadRequestException(failureReason); + } + + await organizationService.AutoAddSeatsAsync(sponsoringOrganization, newSeatsRequired); + } } try diff --git a/test/Core.Test/OrganizationFeatures/OrganizationSponsorships/FamiliesForEnterprise/CreateSponsorshipCommandTests.cs b/test/Core.Test/OrganizationFeatures/OrganizationSponsorships/FamiliesForEnterprise/CreateSponsorshipCommandTests.cs index f6b6721bd2..7dc6b7360d 100644 --- a/test/Core.Test/OrganizationFeatures/OrganizationSponsorships/FamiliesForEnterprise/CreateSponsorshipCommandTests.cs +++ b/test/Core.Test/OrganizationFeatures/OrganizationSponsorships/FamiliesForEnterprise/CreateSponsorshipCommandTests.cs @@ -168,6 +168,11 @@ public class CreateSponsorshipCommandTests : FamiliesForEnterpriseTestsBase }); sutProvider.GetDependency().UserId.Returns(sponsoringOrgUser.UserId.Value); + // Setup for checking available seats + sutProvider.GetDependency() + .GetOccupiedSeatCountByOrganizationIdAsync(sponsoringOrg.Id) + .Returns(0); + await sutProvider.Sut.CreateSponsorshipAsync(sponsoringOrg, sponsoringOrgUser, PlanSponsorshipType.FamiliesForEnterprise, sponsoredEmail, friendlyName, false, null); @@ -293,6 +298,7 @@ public class CreateSponsorshipCommandTests : FamiliesForEnterpriseTestsBase { sponsoringOrg.PlanType = PlanType.EnterpriseAnnually; sponsoringOrg.UseAdminSponsoredFamilies = true; + sponsoringOrg.Seats = 10; sponsoringOrgUser.Status = OrganizationUserStatusType.Confirmed; sutProvider.GetDependency().GetUserByIdAsync(sponsoringOrgUser.UserId!.Value).Returns(user); @@ -311,6 +317,11 @@ public class CreateSponsorshipCommandTests : FamiliesForEnterpriseTestsBase } ]); + // Setup for checking available seats - organization has plenty of seats + sutProvider.GetDependency() + .GetOccupiedSeatCountByOrganizationIdAsync(sponsoringOrg.Id) + .Returns(5); + var actual = await sutProvider.Sut.CreateSponsorshipAsync(sponsoringOrg, sponsoringOrgUser, PlanSponsorshipType.FamiliesForEnterprise, sponsoredEmail, friendlyName, true, notes); @@ -331,5 +342,121 @@ public class CreateSponsorshipCommandTests : FamiliesForEnterpriseTestsBase await sutProvider.GetDependency().Received(1) .CreateAsync(Arg.Is(s => SponsorshipValidator(s, expectedSponsorship))); + + // Verify we didn't need to add seats + await sutProvider.GetDependency().DidNotReceive() + .AutoAddSeatsAsync(Arg.Any(), Arg.Any()); + } + + [Theory] + [BitAutoData(OrganizationUserType.Admin)] + [BitAutoData(OrganizationUserType.Owner)] + public async Task CreateSponsorship_CreatesAdminInitiatedSponsorship_AutoscalesWhenNeeded( + OrganizationUserType organizationUserType, + Organization sponsoringOrg, OrganizationUser sponsoringOrgUser, User user, string sponsoredEmail, + string friendlyName, Guid sponsorshipId, Guid currentUserId, string notes, SutProvider sutProvider) + { + sponsoringOrg.PlanType = PlanType.EnterpriseAnnually; + sponsoringOrg.UseAdminSponsoredFamilies = true; + sponsoringOrg.Seats = 10; + sponsoringOrgUser.Status = OrganizationUserStatusType.Confirmed; + + sutProvider.GetDependency().GetUserByIdAsync(sponsoringOrgUser.UserId!.Value).Returns(user); + sutProvider.GetDependency().WhenForAnyArgs(x => x.UpsertAsync(null!)).Do(callInfo => + { + var sponsorship = callInfo.Arg(); + sponsorship.Id = sponsorshipId; + }); + sutProvider.GetDependency().UserId.Returns(currentUserId); + sutProvider.GetDependency().Organizations.Returns([ + new() + { + Id = sponsoringOrg.Id, + Permissions = new Permissions { ManageUsers = true }, + Type = organizationUserType + } + ]); + + // Setup for checking available seats - organization has no available seats + sutProvider.GetDependency() + .GetOccupiedSeatCountByOrganizationIdAsync(sponsoringOrg.Id) + .Returns(10); + + // Setup for checking if can scale + sutProvider.GetDependency() + .CanScaleAsync(sponsoringOrg, 1) + .Returns((true, "")); + + var actual = await sutProvider.Sut.CreateSponsorshipAsync(sponsoringOrg, sponsoringOrgUser, + PlanSponsorshipType.FamiliesForEnterprise, sponsoredEmail, friendlyName, true, notes); + + + var expectedSponsorship = new OrganizationSponsorship + { + Id = sponsorshipId, + SponsoringOrganizationId = sponsoringOrg.Id, + SponsoringOrganizationUserId = sponsoringOrgUser.Id, + FriendlyName = friendlyName, + OfferedToEmail = sponsoredEmail, + PlanSponsorshipType = PlanSponsorshipType.FamiliesForEnterprise, + IsAdminInitiated = true, + Notes = notes + }; + + Assert.True(SponsorshipValidator(expectedSponsorship, actual)); + + await sutProvider.GetDependency().Received(1) + .CreateAsync(Arg.Is(s => SponsorshipValidator(s, expectedSponsorship))); + + // Verify we needed to add seats + await sutProvider.GetDependency().Received(1) + .AutoAddSeatsAsync(sponsoringOrg, 1); + } + + [Theory] + [BitAutoData(OrganizationUserType.Admin)] + [BitAutoData(OrganizationUserType.Owner)] + public async Task CreateSponsorship_CreatesAdminInitiatedSponsorship_ThrowsWhenCannotAutoscale( + OrganizationUserType organizationUserType, + Organization sponsoringOrg, OrganizationUser sponsoringOrgUser, User user, string sponsoredEmail, + string friendlyName, Guid sponsorshipId, Guid currentUserId, string notes, SutProvider sutProvider) + { + sponsoringOrg.PlanType = PlanType.EnterpriseAnnually; + sponsoringOrg.UseAdminSponsoredFamilies = true; + sponsoringOrg.Seats = 10; + sponsoringOrgUser.Status = OrganizationUserStatusType.Confirmed; + + sutProvider.GetDependency().GetUserByIdAsync(sponsoringOrgUser.UserId!.Value).Returns(user); + sutProvider.GetDependency().WhenForAnyArgs(x => x.UpsertAsync(null!)).Do(callInfo => + { + var sponsorship = callInfo.Arg(); + sponsorship.Id = sponsorshipId; + }); + sutProvider.GetDependency().UserId.Returns(currentUserId); + sutProvider.GetDependency().Organizations.Returns([ + new() + { + Id = sponsoringOrg.Id, + Permissions = new Permissions { ManageUsers = true }, + Type = organizationUserType + } + ]); + + // Setup for checking available seats - organization has no available seats + sutProvider.GetDependency() + .GetOccupiedSeatCountByOrganizationIdAsync(sponsoringOrg.Id) + .Returns(10); + + // Setup for checking if can scale - cannot scale + var failureReason = "Seat limit has been reached."; + sutProvider.GetDependency() + .CanScaleAsync(sponsoringOrg, 1) + .Returns((false, failureReason)); + + var exception = await Assert.ThrowsAsync(() => + sutProvider.Sut.CreateSponsorshipAsync(sponsoringOrg, sponsoringOrgUser, + PlanSponsorshipType.FamiliesForEnterprise, sponsoredEmail, friendlyName, true, notes)); + + Assert.Equal(failureReason, exception.Message); } }