[PM 29610]Update Account Storage Endpoint (#6750)

* update account storage endpoint

* Fix the failing test

* Added flag and refactor base on pr comments

* fix the lint error

* Resolve the pr comments

* Fix the failing test

* Fix the failing test

* Return none

* Resolve the lint error

* Fix the failing test

* Add the missing test

* Formatting issues fixed

src/Api/Billing/Controllers/VNext/AccountBillingVNextController.cs

@@ -1,6 +1,7 @@
 using Bit.Api.Billing.Attributes;
 using Bit.Api.Billing.Models.Requests.Payment;
 using Bit.Api.Billing.Models.Requests.Premium;
+using Bit.Api.Billing.Models.Requests.Storage;
 using Bit.Core;
 using Bit.Core.Billing.Licenses.Queries;
 using Bit.Core.Billing.Payment.Commands;
@@ -23,7 +24,8 @@ public class AccountBillingVNextController(
     IGetCreditQuery getCreditQuery,
     IGetPaymentMethodQuery getPaymentMethodQuery,
     IGetUserLicenseQuery getUserLicenseQuery,
-    IUpdatePaymentMethodCommand updatePaymentMethodCommand) : BaseBillingController
+    IUpdatePaymentMethodCommand updatePaymentMethodCommand,
+    IUpdatePremiumStorageCommand updatePremiumStorageCommand) : BaseBillingController
 {
     [HttpGet("credit")]
     [InjectUser]
@@ -88,4 +90,15 @@ public class AccountBillingVNextController(
         var response = await getUserLicenseQuery.Run(user);
         return TypedResults.Ok(response);
     }
+
+    [HttpPut("storage")]
+    [RequireFeature(FeatureFlagKeys.PM29594_UpdateIndividualSubscriptionPage)]
+    [InjectUser]
+    public async Task<IResult> UpdateStorageAsync(
+        [BindNever] User user,
+        [FromBody] StorageUpdateRequest request)
+    {
+        var result = await updatePremiumStorageCommand.Run(user, request.AdditionalStorageGb);
+        return Handle(result);
+    }
 }

src/Api/Billing/Models/Requests/Storage/StorageUpdateRequest.cs

@@ -0,0 +1,35 @@
+using System.ComponentModel.DataAnnotations;
+
+namespace Bit.Api.Billing.Models.Requests.Storage;
+
+/// <summary>
+/// Request model for updating storage allocation on a user's premium subscription.
+/// Allows for both increasing and decreasing storage in an idempotent manner.
+/// </summary>
+public class StorageUpdateRequest : IValidatableObject
+{
+    /// <summary>
+    /// The additional storage in GB beyond the base storage.
+    /// Must be between 0 and the maximum allowed (minus base storage).
+    /// </summary>
+    [Required]
+    [Range(0, 99)]
+    public short AdditionalStorageGb { get; set; }
+
+    public IEnumerable<ValidationResult> Validate(ValidationContext validationContext)
+    {
+        if (AdditionalStorageGb < 0)
+        {
+            yield return new ValidationResult(
+                "Additional storage cannot be negative.",
+                new[] { nameof(AdditionalStorageGb) });
+        }
+
+        if (AdditionalStorageGb > 99)
+        {
+            yield return new ValidationResult(
+                "Maximum additional storage is 99 GB.",
+                new[] { nameof(AdditionalStorageGb) });
+        }
+    }
+}

src/Core/Billing/Extensions/ServiceCollectionExtensions.cs

@@ -53,6 +53,7 @@ public static class ServiceCollectionExtensions
         services.AddScoped<ICreatePremiumCloudHostedSubscriptionCommand, CreatePremiumCloudHostedSubscriptionCommand>();
         services.AddScoped<ICreatePremiumSelfHostedSubscriptionCommand, CreatePremiumSelfHostedSubscriptionCommand>();
         services.AddTransient<IPreviewPremiumTaxCommand, PreviewPremiumTaxCommand>();
+        services.AddScoped<IUpdatePremiumStorageCommand, UpdatePremiumStorageCommand>();
     }
 
     private static void AddPremiumQueries(this IServiceCollection services)

src/Core/Billing/Premium/Commands/UpdatePremiumStorageCommand.cs

@@ -0,0 +1,144 @@
+using Bit.Core.Billing.Commands;
+using Bit.Core.Billing.Pricing;
+using Bit.Core.Billing.Services;
+using Bit.Core.Entities;
+using Bit.Core.Services;
+using Bit.Core.Utilities;
+using Microsoft.Extensions.Logging;
+using OneOf.Types;
+using Stripe;
+
+namespace Bit.Core.Billing.Premium.Commands;
+
+/// <summary>
+/// Updates the storage allocation for a premium user's subscription.
+/// Handles both increases and decreases in storage in an idempotent manner.
+/// </summary>
+public interface IUpdatePremiumStorageCommand
+{
+    /// <summary>
+    /// Updates the user's storage by the specified additional amount.
+    /// </summary>
+    /// <param name="user">The premium user whose storage should be updated.</param>
+    /// <param name="additionalStorageGb">The additional storage amount in GB beyond base storage.</param>
+    /// <returns>A billing command result indicating success or failure.</returns>
+    Task<BillingCommandResult<None>> Run(User user, short additionalStorageGb);
+}
+
+public class UpdatePremiumStorageCommand(
+    IStripeAdapter stripeAdapter,
+    IUserService userService,
+    IPricingClient pricingClient,
+    ILogger<UpdatePremiumStorageCommand> logger)
+    : BaseBillingCommand<UpdatePremiumStorageCommand>(logger), IUpdatePremiumStorageCommand
+{
+    public Task<BillingCommandResult<None>> Run(User user, short additionalStorageGb) => HandleAsync<None>(async () =>
+    {
+        if (!user.Premium)
+        {
+            return new BadRequest("User does not have a premium subscription.");
+        }
+
+        if (!user.MaxStorageGb.HasValue)
+        {
+            return new BadRequest("No access to storage.");
+        }
+
+        // Fetch all premium plans and the user's subscription to find which plan they're on
+        var premiumPlans = await pricingClient.ListPremiumPlans();
+        var subscription = await stripeAdapter.GetSubscriptionAsync(user.GatewaySubscriptionId);
+
+        // Find the password manager subscription item (seat, not storage) and match it to a plan
+        var passwordManagerItem = subscription.Items.Data.FirstOrDefault(i =>
+            premiumPlans.Any(p => p.Seat.StripePriceId == i.Price.Id));
+
+        if (passwordManagerItem == null)
+        {
+            return new BadRequest("Premium subscription item not found.");
+        }
+
+        var premiumPlan = premiumPlans.First(p => p.Seat.StripePriceId == passwordManagerItem.Price.Id);
+
+        var baseStorageGb = (short)premiumPlan.Storage.Provided;
+
+        if (additionalStorageGb < 0)
+        {
+            return new BadRequest("Additional storage cannot be negative.");
+        }
+
+        var newTotalStorageGb = (short)(baseStorageGb + additionalStorageGb);
+
+        if (newTotalStorageGb > 100)
+        {
+            return new BadRequest("Maximum storage is 100 GB.");
+        }
+
+        // Idempotency check: if user already has the requested storage, return success
+        if (user.MaxStorageGb == newTotalStorageGb)
+        {
+            return new None();
+        }
+
+        var remainingStorage = user.StorageBytesRemaining(newTotalStorageGb);
+        if (remainingStorage < 0)
+        {
+            return new BadRequest(
+                $"You are currently using {CoreHelpers.ReadableBytesSize(user.Storage.GetValueOrDefault(0))} of storage. " +
+                "Delete some stored data first.");
+        }
+
+        // Find the storage line item in the subscription
+        var storageItem = subscription.Items.Data.FirstOrDefault(i => i.Price.Id == premiumPlan.Storage.StripePriceId);
+
+        var subscriptionItemOptions = new List<SubscriptionItemOptions>();
+
+        if (additionalStorageGb > 0)
+        {
+            if (storageItem != null)
+            {
+                // Update existing storage item
+                subscriptionItemOptions.Add(new SubscriptionItemOptions
+                {
+                    Id = storageItem.Id,
+                    Price = premiumPlan.Storage.StripePriceId,
+                    Quantity = additionalStorageGb
+                });
+            }
+            else
+            {
+                // Add new storage item
+                subscriptionItemOptions.Add(new SubscriptionItemOptions
+                {
+                    Price = premiumPlan.Storage.StripePriceId,
+                    Quantity = additionalStorageGb
+                });
+            }
+        }
+        else if (storageItem != null)
+        {
+            // Remove storage item if setting to 0
+            subscriptionItemOptions.Add(new SubscriptionItemOptions
+            {
+                Id = storageItem.Id,
+                Deleted = true
+            });
+        }
+
+        // Update subscription with prorations
+        // Storage is billed annually, so we create prorations and invoice immediately
+        var subscriptionUpdateOptions = new SubscriptionUpdateOptions
+        {
+            Items = subscriptionItemOptions,
+            ProrationBehavior = Core.Constants.CreateProrations
+        };
+
+        await stripeAdapter.UpdateSubscriptionAsync(subscription.Id, subscriptionUpdateOptions);
+
+        // Update the user's max storage
+        user.MaxStorageGb = newTotalStorageGb;
+        await userService.SaveUserAsync(user);
+
+        // No payment intent needed - the subscription update will automatically create and finalize the invoice
+        return new None();
+    });
+}

src/Core/Constants.cs

@@ -196,6 +196,7 @@ public static class FeatureFlagKeys
     public const string PM26462_Milestone_3 = "pm-26462-milestone-3";
     public const string PM28265_EnableReconcileAdditionalStorageJob = "pm-28265-enable-reconcile-additional-storage-job";
     public const string PM28265_ReconcileAdditionalStorageJobEnableLiveMode = "pm-28265-reconcile-additional-storage-job-enable-live-mode";
+    public const string PM29594_UpdateIndividualSubscriptionPage = "pm-29594-update-individual-subscription-page";
 
     /* Key Management Team */
     public const string PrivateKeyRegeneration = "pm-12241-private-key-regeneration";

test/Api.Test/Billing/Controllers/VNext/AccountBillingVNextControllerTests.cs

@@ -1,57 +1,245 @@
 using Bit.Api.Billing.Controllers.VNext;
+using Bit.Api.Billing.Models.Requests.Storage;
+using Bit.Core.Billing.Commands;
 using Bit.Core.Billing.Licenses.Queries;
-using Bit.Core.Billing.Payment.Commands;
-using Bit.Core.Billing.Payment.Queries;
 using Bit.Core.Billing.Premium.Commands;
 using Bit.Core.Entities;
 using Bit.Test.Common.AutoFixture.Attributes;
 using Microsoft.AspNetCore.Http;
 using NSubstitute;
+using OneOf.Types;
 using Xunit;
+using BadRequest = Bit.Core.Billing.Commands.BadRequest;
 
 namespace Bit.Api.Test.Billing.Controllers.VNext;
 
 public class AccountBillingVNextControllerTests
 {
-    private readonly ICreateBitPayInvoiceForCreditCommand _createBitPayInvoiceForCreditCommand;
-    private readonly ICreatePremiumCloudHostedSubscriptionCommand _createPremiumCloudHostedSubscriptionCommand;
-    private readonly IGetCreditQuery _getCreditQuery;
-    private readonly IGetPaymentMethodQuery _getPaymentMethodQuery;
+    private readonly IUpdatePremiumStorageCommand _updatePremiumStorageCommand;
     private readonly IGetUserLicenseQuery _getUserLicenseQuery;
-    private readonly IUpdatePaymentMethodCommand _updatePaymentMethodCommand;
     private readonly AccountBillingVNextController _sut;
 
     public AccountBillingVNextControllerTests()
     {
-        _createBitPayInvoiceForCreditCommand = Substitute.For<ICreateBitPayInvoiceForCreditCommand>();
-        _createPremiumCloudHostedSubscriptionCommand = Substitute.For<ICreatePremiumCloudHostedSubscriptionCommand>();
-        _getCreditQuery = Substitute.For<IGetCreditQuery>();
-        _getPaymentMethodQuery = Substitute.For<IGetPaymentMethodQuery>();
+        _updatePremiumStorageCommand = Substitute.For<IUpdatePremiumStorageCommand>();
         _getUserLicenseQuery = Substitute.For<IGetUserLicenseQuery>();
-        _updatePaymentMethodCommand = Substitute.For<IUpdatePaymentMethodCommand>();
 
         _sut = new AccountBillingVNextController(
-            _createBitPayInvoiceForCreditCommand,
-            _createPremiumCloudHostedSubscriptionCommand,
-            _getCreditQuery,
-            _getPaymentMethodQuery,
+            Substitute.For<Core.Billing.Payment.Commands.ICreateBitPayInvoiceForCreditCommand>(),
+            Substitute.For<Core.Billing.Premium.Commands.ICreatePremiumCloudHostedSubscriptionCommand>(),
+            Substitute.For<Core.Billing.Payment.Queries.IGetCreditQuery>(),
+            Substitute.For<Core.Billing.Payment.Queries.IGetPaymentMethodQuery>(),
             _getUserLicenseQuery,
-            _updatePaymentMethodCommand);
+            Substitute.For<Core.Billing.Payment.Commands.IUpdatePaymentMethodCommand>(),
+            _updatePremiumStorageCommand);
     }
 
     [Theory, BitAutoData]
-    public async Task GetLicenseAsync_ValidUser_ReturnsLicenseResponse(User user,
+    public async Task GetLicenseAsync_ValidUser_ReturnsLicenseResponse(
+        User user,
         Core.Billing.Licenses.Models.Api.Response.LicenseResponseModel licenseResponse)
     {
         // Arrange
         _getUserLicenseQuery.Run(user).Returns(licenseResponse);
-
         // Act
         var result = await _sut.GetLicenseAsync(user);
-
         // Assert
         var okResult = Assert.IsAssignableFrom<IResult>(result);
         await _getUserLicenseQuery.Received(1).Run(user);
     }
 
+    [Theory, BitAutoData]
+    public async Task UpdateStorageAsync_Success_ReturnsOk(User user)
+    {
+        // Arrange
+        var request = new StorageUpdateRequest { AdditionalStorageGb = 10 };
+
+        _updatePremiumStorageCommand.Run(
+            Arg.Is<User>(u => u.Id == user.Id),
+            Arg.Is<short>(s => s == 10))
+            .Returns(new BillingCommandResult<None>(new None()));
+
+        // Act
+        var result = await _sut.UpdateStorageAsync(user, request);
+
+        // Assert
+        var okResult = Assert.IsAssignableFrom<IResult>(result);
+        await _updatePremiumStorageCommand.Received(1).Run(user, 10);
+    }
+
+    [Theory, BitAutoData]
+    public async Task UpdateStorageAsync_UserNotPremium_ReturnsBadRequest(User user)
+    {
+        // Arrange
+        var request = new StorageUpdateRequest { AdditionalStorageGb = 10 };
+        var errorMessage = "User does not have a premium subscription.";
+
+        _updatePremiumStorageCommand.Run(
+            Arg.Is<User>(u => u.Id == user.Id),
+            Arg.Is<short>(s => s == 10))
+            .Returns(new BadRequest(errorMessage));
+
+        // Act
+        var result = await _sut.UpdateStorageAsync(user, request);
+
+        // Assert
+        var badRequestResult = Assert.IsAssignableFrom<IResult>(result);
+        await _updatePremiumStorageCommand.Received(1).Run(user, 10);
+    }
+
+    [Theory, BitAutoData]
+    public async Task UpdateStorageAsync_NoPaymentMethod_ReturnsBadRequest(User user)
+    {
+        // Arrange
+        var request = new StorageUpdateRequest { AdditionalStorageGb = 10 };
+        var errorMessage = "No payment method found.";
+
+        _updatePremiumStorageCommand.Run(
+            Arg.Is<User>(u => u.Id == user.Id),
+            Arg.Is<short>(s => s == 10))
+            .Returns(new BadRequest(errorMessage));
+
+        // Act
+        var result = await _sut.UpdateStorageAsync(user, request);
+
+        // Assert
+        var badRequestResult = Assert.IsAssignableFrom<IResult>(result);
+        await _updatePremiumStorageCommand.Received(1).Run(user, 10);
+    }
+
+    [Theory, BitAutoData]
+    public async Task UpdateStorageAsync_StorageLessThanBase_ReturnsBadRequest(User user)
+    {
+        // Arrange
+        var request = new StorageUpdateRequest { AdditionalStorageGb = 1 };
+        var errorMessage = "Storage cannot be less than the base amount of 1 GB.";
+
+        _updatePremiumStorageCommand.Run(
+            Arg.Is<User>(u => u.Id == user.Id),
+            Arg.Is<short>(s => s == 1))
+            .Returns(new BadRequest(errorMessage));
+
+        // Act
+        var result = await _sut.UpdateStorageAsync(user, request);
+
+        // Assert
+        var badRequestResult = Assert.IsAssignableFrom<IResult>(result);
+        await _updatePremiumStorageCommand.Received(1).Run(user, 1);
+    }
+
+    [Theory, BitAutoData]
+    public async Task UpdateStorageAsync_StorageExceedsMaximum_ReturnsBadRequest(User user)
+    {
+        // Arrange
+        var request = new StorageUpdateRequest { AdditionalStorageGb = 100 };
+        var errorMessage = "Maximum storage is 100 GB.";
+
+        _updatePremiumStorageCommand.Run(
+            Arg.Is<User>(u => u.Id == user.Id),
+            Arg.Is<short>(s => s == 100))
+            .Returns(new BadRequest(errorMessage));
+
+        // Act
+        var result = await _sut.UpdateStorageAsync(user, request);
+
+        // Assert
+        var badRequestResult = Assert.IsAssignableFrom<IResult>(result);
+        await _updatePremiumStorageCommand.Received(1).Run(user, 100);
+    }
+
+    [Theory, BitAutoData]
+    public async Task UpdateStorageAsync_StorageExceedsCurrentUsage_ReturnsBadRequest(User user)
+    {
+        // Arrange
+        var request = new StorageUpdateRequest { AdditionalStorageGb = 2 };
+        var errorMessage = "You are currently using 5.00 GB of storage. Delete some stored data first.";
+
+        _updatePremiumStorageCommand.Run(
+            Arg.Is<User>(u => u.Id == user.Id),
+            Arg.Is<short>(s => s == 2))
+            .Returns(new BadRequest(errorMessage));
+
+        // Act
+        var result = await _sut.UpdateStorageAsync(user, request);
+
+        // Assert
+        var badRequestResult = Assert.IsAssignableFrom<IResult>(result);
+        await _updatePremiumStorageCommand.Received(1).Run(user, 2);
+    }
+
+    [Theory, BitAutoData]
+    public async Task UpdateStorageAsync_IncreaseStorage_Success(User user)
+    {
+        // Arrange
+        var request = new StorageUpdateRequest { AdditionalStorageGb = 15 };
+
+        _updatePremiumStorageCommand.Run(
+            Arg.Is<User>(u => u.Id == user.Id),
+            Arg.Is<short>(s => s == 15))
+            .Returns(new BillingCommandResult<None>(new None()));
+
+        // Act
+        var result = await _sut.UpdateStorageAsync(user, request);
+
+        // Assert
+        var okResult = Assert.IsAssignableFrom<IResult>(result);
+        await _updatePremiumStorageCommand.Received(1).Run(user, 15);
+    }
+
+    [Theory, BitAutoData]
+    public async Task UpdateStorageAsync_DecreaseStorage_Success(User user)
+    {
+        // Arrange
+        var request = new StorageUpdateRequest { AdditionalStorageGb = 3 };
+
+        _updatePremiumStorageCommand.Run(
+            Arg.Is<User>(u => u.Id == user.Id),
+            Arg.Is<short>(s => s == 3))
+            .Returns(new BillingCommandResult<None>(new None()));
+
+        // Act
+        var result = await _sut.UpdateStorageAsync(user, request);
+
+        // Assert
+        var okResult = Assert.IsAssignableFrom<IResult>(result);
+        await _updatePremiumStorageCommand.Received(1).Run(user, 3);
+    }
+
+    [Theory, BitAutoData]
+    public async Task UpdateStorageAsync_MaximumStorage_Success(User user)
+    {
+        // Arrange
+        var request = new StorageUpdateRequest { AdditionalStorageGb = 100 };
+
+        _updatePremiumStorageCommand.Run(
+            Arg.Is<User>(u => u.Id == user.Id),
+            Arg.Is<short>(s => s == 100))
+            .Returns(new BillingCommandResult<None>(new None()));
+
+        // Act
+        var result = await _sut.UpdateStorageAsync(user, request);
+
+        // Assert
+        var okResult = Assert.IsAssignableFrom<IResult>(result);
+        await _updatePremiumStorageCommand.Received(1).Run(user, 100);
+    }
+
+    [Theory, BitAutoData]
+    public async Task UpdateStorageAsync_NullPaymentSecret_Success(User user)
+    {
+        // Arrange
+        var request = new StorageUpdateRequest { AdditionalStorageGb = 5 };
+
+        _updatePremiumStorageCommand.Run(
+            Arg.Is<User>(u => u.Id == user.Id),
+            Arg.Is<short>(s => s == 5))
+            .Returns(new BillingCommandResult<None>(new None()));
+
+        // Act
+        var result = await _sut.UpdateStorageAsync(user, request);
+
+        // Assert
+        var okResult = Assert.IsAssignableFrom<IResult>(result);
+        await _updatePremiumStorageCommand.Received(1).Run(user, 5);
+    }
 }

test/Core.Test/Billing/Payment/Commands/UpdateBillingAddressCommandTests.cs

@@ -407,4 +407,85 @@ public class UpdateBillingAddressCommandTests
             options => options.Type == TaxIdType.SpanishNIF &&
                        options.Value == input.TaxId.Value));
     }
+
+    [Fact]
+    public async Task Run_BusinessOrganization_UpdatingWithSameTaxId_DeletesBeforeCreating()
+    {
+        var organization = new Organization
+        {
+            PlanType = PlanType.EnterpriseAnnually,
+            GatewayCustomerId = "cus_123",
+            GatewaySubscriptionId = "sub_123"
+        };
+
+        var input = new BillingAddress
+        {
+            Country = "US",
+            PostalCode = "12345",
+            Line1 = "123 Main St.",
+            Line2 = "Suite 100",
+            City = "New York",
+            State = "NY",
+            TaxId = new TaxID("us_ein", "987654321")
+        };
+
+        var existingTaxId = new TaxId { Id = "tax_id_123", Type = "us_ein", Value = "987654321" };
+
+        var customer = new Customer
+        {
+            Address = new Address
+            {
+                Country = "US",
+                PostalCode = "12345",
+                Line1 = "123 Main St.",
+                Line2 = "Suite 100",
+                City = "New York",
+                State = "NY"
+            },
+            Id = organization.GatewayCustomerId,
+            Subscriptions = new StripeList<Subscription>
+            {
+                Data =
+                [
+                    new Subscription
+                    {
+                        Id = organization.GatewaySubscriptionId,
+                        AutomaticTax = new SubscriptionAutomaticTax { Enabled = false }
+                    }
+                ]
+            },
+            TaxIds = new StripeList<TaxId>
+            {
+                Data = [existingTaxId]
+            }
+        };
+
+        _stripeAdapter.UpdateCustomerAsync(organization.GatewayCustomerId, Arg.Is<CustomerUpdateOptions>(options =>
+            options.Address.Matches(input) &&
+            options.HasExpansions("subscriptions", "tax_ids") &&
+            options.TaxExempt == TaxExempt.None
+        )).Returns(customer);
+
+        var newTaxId = new TaxId { Id = "tax_id_456", Type = "us_ein", Value = "987654321" };
+        _stripeAdapter.CreateTaxIdAsync(customer.Id, Arg.Is<TaxIdCreateOptions>(
+            options => options.Type == "us_ein" && options.Value == "987654321"
+        )).Returns(newTaxId);
+
+        var result = await _command.Run(organization, input);
+
+        Assert.True(result.IsT0);
+        var output = result.AsT0;
+        Assert.Equivalent(input, output);
+
+        // Verify that deletion happens before creation
+        Received.InOrder(() =>
+        {
+            _stripeAdapter.DeleteTaxIdAsync(customer.Id, existingTaxId.Id);
+            _stripeAdapter.CreateTaxIdAsync(customer.Id, Arg.Any<TaxIdCreateOptions>());
+        });
+
+        await _stripeAdapter.Received(1).DeleteTaxIdAsync(customer.Id, existingTaxId.Id);
+        await _stripeAdapter.Received(1).CreateTaxIdAsync(customer.Id, Arg.Is<TaxIdCreateOptions>(
+            options => options.Type == "us_ein" && options.Value == "987654321"));
+    }
 }

test/Core.Test/Billing/Premium/Commands/UpdatePremiumStorageCommandTests.cs

@@ -0,0 +1,339 @@
+using Bit.Core.Billing.Premium.Commands;
+using Bit.Core.Billing.Pricing;
+using Bit.Core.Billing.Services;
+using Bit.Core.Entities;
+using Bit.Core.Services;
+using Bit.Test.Common.AutoFixture.Attributes;
+using Microsoft.Extensions.Logging;
+using NSubstitute;
+using Stripe;
+using Xunit;
+using PremiumPlan = Bit.Core.Billing.Pricing.Premium.Plan;
+using PremiumPurchasable = Bit.Core.Billing.Pricing.Premium.Purchasable;
+
+namespace Bit.Core.Test.Billing.Premium.Commands;
+
+public class UpdatePremiumStorageCommandTests
+{
+    private readonly IStripeAdapter _stripeAdapter = Substitute.For<IStripeAdapter>();
+    private readonly IUserService _userService = Substitute.For<IUserService>();
+    private readonly IPricingClient _pricingClient = Substitute.For<IPricingClient>();
+    private readonly PremiumPlan _premiumPlan;
+    private readonly UpdatePremiumStorageCommand _command;
+
+    public UpdatePremiumStorageCommandTests()
+    {
+        // Setup default premium plan with standard pricing
+        _premiumPlan = new PremiumPlan
+        {
+            Name = "Premium",
+            Available = true,
+            LegacyYear = null,
+            Seat = new PremiumPurchasable { Price = 10M, StripePriceId = "price_premium", Provided = 1 },
+            Storage = new PremiumPurchasable { Price = 4M, StripePriceId = "price_storage", Provided = 1 }
+        };
+        _pricingClient.ListPremiumPlans().Returns(new List<PremiumPlan> { _premiumPlan });
+
+        _command = new UpdatePremiumStorageCommand(
+            _stripeAdapter,
+            _userService,
+            _pricingClient,
+            Substitute.For<ILogger<UpdatePremiumStorageCommand>>());
+    }
+
+    private Subscription CreateMockSubscription(string subscriptionId, int? storageQuantity = null)
+    {
+        var items = new List<SubscriptionItem>();
+
+        // Always add the seat item
+        items.Add(new SubscriptionItem
+        {
+            Id = "si_seat",
+            Price = new Price { Id = "price_premium" },
+            Quantity = 1
+        });
+
+        // Add storage item if quantity is provided
+        if (storageQuantity.HasValue && storageQuantity.Value > 0)
+        {
+            items.Add(new SubscriptionItem
+            {
+                Id = "si_storage",
+                Price = new Price { Id = "price_storage" },
+                Quantity = storageQuantity.Value
+            });
+        }
+
+        return new Subscription
+        {
+            Id = subscriptionId,
+            Items = new StripeList<SubscriptionItem>
+            {
+                Data = items
+            }
+        };
+    }
+
+    [Theory, BitAutoData]
+    public async Task Run_UserNotPremium_ReturnsBadRequest(User user)
+    {
+        // Arrange
+        user.Premium = false;
+
+        // Act
+        var result = await _command.Run(user, 5);
+
+        // Assert
+        Assert.True(result.IsT1);
+        var badRequest = result.AsT1;
+        Assert.Equal("User does not have a premium subscription.", badRequest.Response);
+    }
+
+    [Theory, BitAutoData]
+    public async Task Run_NegativeStorage_ReturnsBadRequest(User user)
+    {
+        // Arrange
+        user.Premium = true;
+        user.MaxStorageGb = 5;
+        user.GatewaySubscriptionId = "sub_123";
+
+        var subscription = CreateMockSubscription("sub_123", 4);
+        _stripeAdapter.GetSubscriptionAsync("sub_123").Returns(subscription);
+
+        // Act
+        var result = await _command.Run(user, -5);
+
+        // Assert
+        Assert.True(result.IsT1);
+        var badRequest = result.AsT1;
+        Assert.Equal("Additional storage cannot be negative.", badRequest.Response);
+    }
+
+    [Theory, BitAutoData]
+    public async Task Run_StorageExceedsMaximum_ReturnsBadRequest(User user)
+    {
+        // Arrange
+        user.Premium = true;
+        user.MaxStorageGb = 5;
+        user.GatewaySubscriptionId = "sub_123";
+
+        var subscription = CreateMockSubscription("sub_123", 4);
+        _stripeAdapter.GetSubscriptionAsync("sub_123").Returns(subscription);
+
+        // Act
+        var result = await _command.Run(user, 100);
+
+        // Assert
+        Assert.True(result.IsT1);
+        var badRequest = result.AsT1;
+        Assert.Equal("Maximum storage is 100 GB.", badRequest.Response);
+    }
+
+    [Theory, BitAutoData]
+    public async Task Run_NoMaxStorageGb_ReturnsBadRequest(User user)
+    {
+        // Arrange
+        user.Premium = true;
+        user.MaxStorageGb = null;
+
+        // Act
+        var result = await _command.Run(user, 5);
+
+        // Assert
+        Assert.True(result.IsT1);
+        var badRequest = result.AsT1;
+        Assert.Equal("No access to storage.", badRequest.Response);
+    }
+
+    [Theory, BitAutoData]
+    public async Task Run_StorageExceedsCurrentUsage_ReturnsBadRequest(User user)
+    {
+        // Arrange
+        user.Premium = true;
+        user.MaxStorageGb = 10;
+        user.Storage = 5L * 1024 * 1024 * 1024; // 5 GB currently used
+        user.GatewaySubscriptionId = "sub_123";
+
+        var subscription = CreateMockSubscription("sub_123", 9);
+        _stripeAdapter.GetSubscriptionAsync("sub_123").Returns(subscription);
+
+        // Act
+        var result = await _command.Run(user, 0);
+
+        // Assert
+        Assert.True(result.IsT1);
+        var badRequest = result.AsT1;
+        Assert.Contains("You are currently using", badRequest.Response);
+        Assert.Contains("Delete some stored data first", badRequest.Response);
+    }
+
+    [Theory, BitAutoData]
+    public async Task Run_SameStorageAmount_Idempotent(User user)
+    {
+        // Arrange
+        user.Premium = true;
+        user.MaxStorageGb = 5;
+        user.Storage = 2L * 1024 * 1024 * 1024;
+        user.GatewaySubscriptionId = "sub_123";
+
+        var subscription = CreateMockSubscription("sub_123", 4);
+        _stripeAdapter.GetSubscriptionAsync("sub_123").Returns(subscription);
+
+        // Act
+        var result = await _command.Run(user, 4);
+
+        // Assert
+        Assert.True(result.IsT0);
+
+        // Verify subscription was fetched but NOT updated
+        await _stripeAdapter.Received(1).GetSubscriptionAsync("sub_123");
+        await _stripeAdapter.DidNotReceive().UpdateSubscriptionAsync(Arg.Any<string>(), Arg.Any<SubscriptionUpdateOptions>());
+        await _userService.DidNotReceive().SaveUserAsync(Arg.Any<User>());
+    }
+
+    [Theory, BitAutoData]
+    public async Task Run_IncreaseStorage_Success(User user)
+    {
+        // Arrange
+        user.Premium = true;
+        user.MaxStorageGb = 5;
+        user.Storage = 2L * 1024 * 1024 * 1024;
+        user.GatewaySubscriptionId = "sub_123";
+
+        var subscription = CreateMockSubscription("sub_123", 4);
+        _stripeAdapter.GetSubscriptionAsync("sub_123").Returns(subscription);
+
+        // Act
+        var result = await _command.Run(user, 9);
+
+        // Assert
+        Assert.True(result.IsT0);
+
+        // Verify subscription was updated
+        await _stripeAdapter.Received(1).UpdateSubscriptionAsync(
+            "sub_123",
+            Arg.Is<SubscriptionUpdateOptions>(opts =>
+                opts.Items.Count == 1 &&
+                opts.Items[0].Id == "si_storage" &&
+                opts.Items[0].Quantity == 9 &&
+                opts.ProrationBehavior == "create_prorations"));
+
+        // Verify user was saved
+        await _userService.Received(1).SaveUserAsync(Arg.Is<User>(u =>
+            u.Id == user.Id &&
+            u.MaxStorageGb == 10));
+    }
+
+    [Theory, BitAutoData]
+    public async Task Run_AddStorageFromZero_Success(User user)
+    {
+        // Arrange
+        user.Premium = true;
+        user.MaxStorageGb = 1;
+        user.Storage = 500L * 1024 * 1024;
+        user.GatewaySubscriptionId = "sub_123";
+
+        var subscription = CreateMockSubscription("sub_123", null);
+        _stripeAdapter.GetSubscriptionAsync("sub_123").Returns(subscription);
+
+        // Act
+        var result = await _command.Run(user, 9);
+
+        // Assert
+        Assert.True(result.IsT0);
+
+        // Verify subscription was updated with new storage item
+        await _stripeAdapter.Received(1).UpdateSubscriptionAsync(
+            "sub_123",
+            Arg.Is<SubscriptionUpdateOptions>(opts =>
+                opts.Items.Count == 1 &&
+                opts.Items[0].Price == "price_storage" &&
+                opts.Items[0].Quantity == 9));
+
+        await _userService.Received(1).SaveUserAsync(Arg.Is<User>(u => u.MaxStorageGb == 10));
+    }
+
+    [Theory, BitAutoData]
+    public async Task Run_DecreaseStorage_Success(User user)
+    {
+        // Arrange
+        user.Premium = true;
+        user.MaxStorageGb = 10;
+        user.Storage = 2L * 1024 * 1024 * 1024;
+        user.GatewaySubscriptionId = "sub_123";
+
+        var subscription = CreateMockSubscription("sub_123", 9);
+        _stripeAdapter.GetSubscriptionAsync("sub_123").Returns(subscription);
+
+        // Act
+        var result = await _command.Run(user, 2);
+
+        // Assert
+        Assert.True(result.IsT0);
+
+        // Verify subscription was updated
+        await _stripeAdapter.Received(1).UpdateSubscriptionAsync(
+            "sub_123",
+            Arg.Is<SubscriptionUpdateOptions>(opts =>
+                opts.Items.Count == 1 &&
+                opts.Items[0].Id == "si_storage" &&
+                opts.Items[0].Quantity == 2));
+
+        await _userService.Received(1).SaveUserAsync(Arg.Is<User>(u => u.MaxStorageGb == 3));
+    }
+
+    [Theory, BitAutoData]
+    public async Task Run_RemoveAllAdditionalStorage_Success(User user)
+    {
+        // Arrange
+        user.Premium = true;
+        user.MaxStorageGb = 10;
+        user.Storage = 500L * 1024 * 1024;
+        user.GatewaySubscriptionId = "sub_123";
+
+        var subscription = CreateMockSubscription("sub_123", 9);
+        _stripeAdapter.GetSubscriptionAsync("sub_123").Returns(subscription);
+
+        // Act
+        var result = await _command.Run(user, 0);
+
+        // Assert
+        Assert.True(result.IsT0);
+
+        // Verify subscription item was deleted
+        await _stripeAdapter.Received(1).UpdateSubscriptionAsync(
+            "sub_123",
+            Arg.Is<SubscriptionUpdateOptions>(opts =>
+                opts.Items.Count == 1 &&
+                opts.Items[0].Id == "si_storage" &&
+                opts.Items[0].Deleted == true));
+
+        await _userService.Received(1).SaveUserAsync(Arg.Is<User>(u => u.MaxStorageGb == 1));
+    }
+
+    [Theory, BitAutoData]
+    public async Task Run_MaximumStorage_Success(User user)
+    {
+        // Arrange
+        user.Premium = true;
+        user.MaxStorageGb = 5;
+        user.Storage = 2L * 1024 * 1024 * 1024;
+        user.GatewaySubscriptionId = "sub_123";
+
+        var subscription = CreateMockSubscription("sub_123", 4);
+        _stripeAdapter.GetSubscriptionAsync("sub_123").Returns(subscription);
+
+        // Act
+        var result = await _command.Run(user, 99);
+
+        // Assert
+        Assert.True(result.IsT0);
+
+        await _stripeAdapter.Received(1).UpdateSubscriptionAsync(
+            "sub_123",
+            Arg.Is<SubscriptionUpdateOptions>(opts =>
+                opts.Items[0].Quantity == 99));
+
+        await _userService.Received(1).SaveUserAsync(Arg.Is<User>(u => u.MaxStorageGb == 100));
+    }
+}