[PM-34866][PM-34865] Fix EnableAutomaticTaxAsync to update schedule phases (#7437)

connerbw · connerbw · commit 51455ecbf20f · 2026-04-10T18:07:31.000Z
* [PM-34866] Fix EnableAutomaticTaxAsync to update schedule phases

* Use test clock frozen time for phase filtering

* Expand test_clock on customer subscription fetches
diff --git a/src/Billing/Services/Implementations/UpcomingInvoiceHandler.cs b/src/Billing/Services/Implementations/UpcomingInvoiceHandler.cs
@@ -50,7 +50,7 @@ public async Task HandleAsync(Event parsedEvent)
 
         var customer =
             await stripeAdapter.GetCustomerAsync(invoice.CustomerId,
-                new CustomerGetOptions { Expand = ["subscriptions", "tax", "tax_ids"] });
+                new CustomerGetOptions { Expand = ["subscriptions", "subscriptions.data.test_clock", "tax", "tax_ids"] });
 
         var subscription = customer.Subscriptions.FirstOrDefault();
 
@@ -574,6 +574,47 @@ private async Task EnableAutomaticTaxAsync(Subscription subscription)
 
             if (activeSchedule != null)
             {
+                var now = subscription.TestClock?.FrozenTime ?? DateTime.UtcNow;
+                var phases = new List<SubscriptionSchedulePhaseOptions>();
+
+                for (var i = 0; i < activeSchedule.Phases.Count; i++)
+                {
+                    var phase = activeSchedule.Phases[i];
+
+                    // Skip phases that have already completed
+                    if (phase.EndDate <= now)
+                    {
+                        continue;
+                    }
+
+                    // When a phase's predecessor has ended, the phase is already active and
+                    // its one-time migration discount has been applied and consumed.
+                    // Re-including it would cause Stripe to re-apply it.
+                    var discountConsumed = i > 0 && activeSchedule.Phases[i - 1].EndDate <= now;
+
+                    phases.Add(new SubscriptionSchedulePhaseOptions
+                    {
+                        StartDate = phase.StartDate,
+                        EndDate = phase.EndDate,
+                        Items = phase.Items.Select(item => new SubscriptionSchedulePhaseItemOptions
+                        {
+                            Price = item.PriceId,
+                            Quantity = item.Quantity
+                        }).ToList(),
+                        Discounts = discountConsumed
+                            ? []
+                            : phase.Discounts?.Select(d => new SubscriptionSchedulePhaseDiscountOptions
+                            {
+                                Coupon = d.CouponId
+                            }).ToList(),
+                        ProrationBehavior = phase.ProrationBehavior,
+                        AutomaticTax = new SubscriptionSchedulePhaseAutomaticTaxOptions
+                        {
+                            Enabled = true
+                        }
+                    });
+                }
+
                 await stripeAdapter.UpdateSubscriptionScheduleAsync(activeSchedule.Id,
                     new SubscriptionScheduleUpdateOptions
                     {
@@ -583,7 +624,8 @@ await stripeAdapter.UpdateSubscriptionScheduleAsync(activeSchedule.Id,
                             {
                                 Enabled = true
                             }
-                        }
+                        },
+                        Phases = phases
                     });
                 return;
             }
diff --git a/src/Core/Billing/Payment/Commands/UpdateBillingAddressCommand.cs b/src/Core/Billing/Payment/Commands/UpdateBillingAddressCommand.cs
@@ -60,7 +60,7 @@ await stripeAdapter.UpdateCustomerAsync(subscriber.GatewayCustomerId,
                         City = billingAddress.City,
                         State = billingAddress.State
                     },
-                    Expand = ["subscriptions"]
+                    Expand = ["subscriptions", "subscriptions.data.test_clock"]
                 });
 
         await EnableAutomaticTaxAsync(subscriber, customer);
@@ -86,7 +86,7 @@ private async Task<BillingCommandResult<BillingAddress>> UpdateBusinessBillingAd
                     City = billingAddress.City,
                     State = billingAddress.State
                 },
-                Expand = ["subscriptions", "tax_ids"],
+                Expand = ["subscriptions", "subscriptions.data.test_clock", "tax_ids"],
                 TaxExempt = determinedTaxExemptStatus
             });
 
@@ -149,6 +149,43 @@ private async Task EnableAutomaticTaxAsync(
 
                     if (activeSchedule != null)
                     {
+                        var now = subscription.TestClock?.FrozenTime ?? DateTime.UtcNow;
+                        var phases = new List<SubscriptionSchedulePhaseOptions>();
+
+                        for (var i = 0; i < activeSchedule.Phases.Count; i++)
+                        {
+                            var phase = activeSchedule.Phases[i];
+
+                            if (phase.EndDate <= now)
+                            {
+                                continue;
+                            }
+
+                            var discountConsumed = i > 0 && activeSchedule.Phases[i - 1].EndDate <= now;
+
+                            phases.Add(new SubscriptionSchedulePhaseOptions
+                            {
+                                StartDate = phase.StartDate,
+                                EndDate = phase.EndDate,
+                                Items = phase.Items.Select(item => new SubscriptionSchedulePhaseItemOptions
+                                {
+                                    Price = item.PriceId,
+                                    Quantity = item.Quantity
+                                }).ToList(),
+                                Discounts = discountConsumed
+                                    ? []
+                                    : phase.Discounts?.Select(d => new SubscriptionSchedulePhaseDiscountOptions
+                                    {
+                                        Coupon = d.CouponId
+                                    }).ToList(),
+                                ProrationBehavior = phase.ProrationBehavior,
+                                AutomaticTax = new SubscriptionSchedulePhaseAutomaticTaxOptions
+                                {
+                                    Enabled = true
+                                }
+                            });
+                        }
+
                         await stripeAdapter.UpdateSubscriptionScheduleAsync(activeSchedule.Id,
                             new SubscriptionScheduleUpdateOptions
                             {
@@ -158,7 +195,8 @@ await stripeAdapter.UpdateSubscriptionScheduleAsync(activeSchedule.Id,
                                     {
                                         Enabled = true
                                     }
-                                }
+                                },
+                                Phases = phases
                             });
                         return;
                     }
diff --git a/test/Billing.Test/Services/UpcomingInvoiceHandlerTests.cs b/test/Billing.Test/Services/UpcomingInvoiceHandlerTests.cs
@@ -3104,7 +3104,7 @@ await _stripeAdapter.DidNotReceive().UpdateSubscriptionAsync(
     #region Schedule-Aware Tax Handling
 
     [Fact]
-    public async Task HandleAsync_WhenOrganizationTaxNotEnabled_FlagOn_SchedulePresent_UpdatesScheduleDefaultSettings()
+    public async Task HandleAsync_WhenOrganizationTaxNotEnabled_FlagOn_SchedulePresent_UpdatesSchedulePhasesAndDefaultSettings()
     {
         // Arrange
         var parsedEvent = new Event { Id = "evt_123", Type = "invoice.upcoming" };
@@ -3125,6 +3125,10 @@ public async Task HandleAsync_WhenOrganizationTaxNotEnabled_FlagOn_SchedulePrese
         };
         var organization = new Organization { Id = _organizationId, PlanType = PlanType.TeamsAnnually, BillingEmail = "test@test.com" };
 
+        var phase1Start = DateTime.UtcNow.AddDays(-10);
+        var phase1End = DateTime.UtcNow.AddDays(5);
+        var phase2End = DateTime.UtcNow.AddDays(370);
+
         _stripeEventService.GetInvoice(parsedEvent).Returns(invoice);
         _stripeAdapter.GetCustomerAsync(invoice.CustomerId, Arg.Any<CustomerGetOptions>()).Returns(customer);
         _stripeEventUtilityService.GetIdsFromMetadata(subscription.Metadata)
@@ -3142,19 +3146,44 @@ public async Task HandleAsync_WhenOrganizationTaxNotEnabled_FlagOn_SchedulePrese
                     {
                         Id = "sub_sched_123",
                         SubscriptionId = "sub_123",
-                        Status = SubscriptionScheduleStatus.Active
+                        Status = SubscriptionScheduleStatus.Active,
+                        Phases = new List<SubscriptionSchedulePhase>
+                        {
+                            new()
+                            {
+                                StartDate = phase1Start,
+                                EndDate = phase1End,
+                                Items = [new SubscriptionSchedulePhaseItem { PriceId = "price_old", Quantity = 1 }],
+                                Discounts = [],
+                                ProrationBehavior = "none"
+                            },
+                            new()
+                            {
+                                StartDate = phase1End,
+                                EndDate = phase2End,
+                                Items = [new SubscriptionSchedulePhaseItem { PriceId = "price_new", Quantity = 1 }],
+                                Discounts = [new SubscriptionSchedulePhaseDiscount { CouponId = "milestone-coupon" }],
+                                ProrationBehavior = "none"
+                            }
+                        }
                     }
                 ]
             });
 
         // Act
         await _sut.HandleAsync(parsedEvent);
 
-        // Assert — schedule's default_settings updated
+        // Assert — schedule updated with phases and default_settings
         await _stripeAdapter.Received(1).UpdateSubscriptionScheduleAsync(
             Arg.Is("sub_sched_123"),
             Arg.Is<SubscriptionScheduleUpdateOptions>(o =>
-                o.DefaultSettings.AutomaticTax.Enabled == true));
+                o.DefaultSettings.AutomaticTax.Enabled == true &&
+                o.Phases.Count == 2 &&
+                o.Phases[0].AutomaticTax.Enabled == true &&
+                o.Phases[0].Items[0].Price == "price_old" &&
+                o.Phases[1].AutomaticTax.Enabled == true &&
+                o.Phases[1].Items[0].Price == "price_new" &&
+                o.Phases[1].Discounts[0].Coupon == "milestone-coupon"));
 
         // Assert — subscription NOT updated directly for tax
         await _stripeAdapter.DidNotReceive().UpdateSubscriptionAsync(
@@ -3208,7 +3237,7 @@ await _stripeAdapter.DidNotReceive().UpdateSubscriptionScheduleAsync(
     }
 
     [Fact]
-    public async Task HandleAsync_WhenPremiumUserTaxNotEnabled_FlagOn_SchedulePresent_UpdatesScheduleDefaultSettings()
+    public async Task HandleAsync_WhenPremiumUserTaxNotEnabled_FlagOn_SchedulePresent_UpdatesSchedulePhasesAndDefaultSettings()
     {
         // Arrange
         var parsedEvent = new Event { Id = "evt_123", Type = "invoice.upcoming" };
@@ -3229,6 +3258,10 @@ public async Task HandleAsync_WhenPremiumUserTaxNotEnabled_FlagOn_SchedulePresen
         };
         var user = new User { Id = _userId, Email = "test@test.com", Premium = true };
 
+        var phase1Start = DateTime.UtcNow.AddDays(-10);
+        var phase1End = DateTime.UtcNow.AddDays(5);
+        var phase2End = DateTime.UtcNow.AddDays(370);
+
         _stripeEventService.GetInvoice(parsedEvent).Returns(invoice);
         _stripeAdapter.GetCustomerAsync(invoice.CustomerId, Arg.Any<CustomerGetOptions>()).Returns(customer);
         _stripeEventUtilityService.GetIdsFromMetadata(subscription.Metadata)
@@ -3245,19 +3278,44 @@ public async Task HandleAsync_WhenPremiumUserTaxNotEnabled_FlagOn_SchedulePresen
                     {
                         Id = "sub_sched_456",
                         SubscriptionId = "sub_123",
-                        Status = SubscriptionScheduleStatus.Active
+                        Status = SubscriptionScheduleStatus.Active,
+                        Phases = new List<SubscriptionSchedulePhase>
+                        {
+                            new()
+                            {
+                                StartDate = phase1Start,
+                                EndDate = phase1End,
+                                Items = [new SubscriptionSchedulePhaseItem { PriceId = "premium-annually", Quantity = 1 }],
+                                Discounts = [],
+                                ProrationBehavior = "none"
+                            },
+                            new()
+                            {
+                                StartDate = phase1End,
+                                EndDate = phase2End,
+                                Items = [new SubscriptionSchedulePhaseItem { PriceId = "premium-annually-new", Quantity = 1 }],
+                                Discounts = [new SubscriptionSchedulePhaseDiscount { CouponId = "milestone-2c" }],
+                                ProrationBehavior = "none"
+                            }
+                        }
                     }
                 ]
             });
 
         // Act
         await _sut.HandleAsync(parsedEvent);
 
-        // Assert — schedule's default_settings updated
+        // Assert — schedule updated with phases and default_settings
         await _stripeAdapter.Received(1).UpdateSubscriptionScheduleAsync(
             Arg.Is("sub_sched_456"),
             Arg.Is<SubscriptionScheduleUpdateOptions>(o =>
-                o.DefaultSettings.AutomaticTax.Enabled == true));
+                o.DefaultSettings.AutomaticTax.Enabled == true &&
+                o.Phases.Count == 2 &&
+                o.Phases[0].AutomaticTax.Enabled == true &&
+                o.Phases[0].Items[0].Price == "premium-annually" &&
+                o.Phases[1].AutomaticTax.Enabled == true &&
+                o.Phases[1].Items[0].Price == "premium-annually-new" &&
+                o.Phases[1].Discounts[0].Coupon == "milestone-2c"));
 
         // Assert — subscription NOT updated directly for tax
         await _stripeAdapter.DidNotReceive().UpdateSubscriptionAsync(
@@ -3309,6 +3367,88 @@ await _stripeAdapter.DidNotReceive().UpdateSubscriptionScheduleAsync(
             Arg.Any<string>(), Arg.Any<SubscriptionScheduleUpdateOptions>());
     }
 
+    [Fact]
+    public async Task HandleAsync_WhenTaxNotEnabled_FlagOn_Phase2Active_SkipsCompletedPhaseAndClearsConsumedDiscounts()
+    {
+        // Arrange — Phase 1 has ended, Phase 2 is now the active phase.
+        // Phase 2's one-time migration discount was consumed at transition and must not be re-included.
+        var parsedEvent = new Event { Id = "evt_123", Type = "invoice.upcoming" };
+        var invoice = new Invoice { CustomerId = "cus_123", Lines = new StripeList<InvoiceLineItem> { Data = [] } };
+        var subscription = new Subscription
+        {
+            Id = "sub_123",
+            CustomerId = "cus_123",
+            AutomaticTax = new SubscriptionAutomaticTax { Enabled = false },
+            Items = new StripeList<SubscriptionItem> { Data = [] },
+            Metadata = new Dictionary<string, string> { { "userId", _userId.ToString() } }
+        };
+        var customer = new Customer
+        {
+            Id = "cus_123",
+            Subscriptions = new StripeList<Subscription> { Data = [subscription] },
+            Tax = new CustomerTax { AutomaticTax = AutomaticTaxStatus.Supported }
+        };
+        var user = new User { Id = _userId, Email = "test@test.com", Premium = true };
+
+        // Phase 1 ended yesterday, Phase 2 active now
+        var phase1Start = DateTime.UtcNow.AddDays(-375);
+        var phase1End = DateTime.UtcNow.AddDays(-1);
+        var phase2End = DateTime.UtcNow.AddDays(364);
+
+        _stripeEventService.GetInvoice(parsedEvent).Returns(invoice);
+        _stripeAdapter.GetCustomerAsync(invoice.CustomerId, Arg.Any<CustomerGetOptions>()).Returns(customer);
+        _stripeEventUtilityService.GetIdsFromMetadata(subscription.Metadata)
+            .Returns(new Tuple<Guid?, Guid?, Guid?>(null, _userId, null));
+        _userRepository.GetByIdAsync(_userId).Returns(user);
+        _featureService.IsEnabled(FeatureFlagKeys.PM32645_DeferPriceMigrationToRenewal).Returns(true);
+
+        _stripeAdapter.ListSubscriptionSchedulesAsync(Arg.Any<SubscriptionScheduleListOptions>())
+            .Returns(new StripeList<SubscriptionSchedule>
+            {
+                Data =
+                [
+                    new SubscriptionSchedule
+                    {
+                        Id = "sub_sched_789",
+                        SubscriptionId = "sub_123",
+                        Status = SubscriptionScheduleStatus.Active,
+                        Phases = new List<SubscriptionSchedulePhase>
+                        {
+                            new()
+                            {
+                                StartDate = phase1Start,
+                                EndDate = phase1End,
+                                Items = [new SubscriptionSchedulePhaseItem { PriceId = "price_old", Quantity = 1 }],
+                                Discounts = [],
+                                ProrationBehavior = "none"
+                            },
+                            new()
+                            {
+                                StartDate = phase1End,
+                                EndDate = phase2End,
+                                Items = [new SubscriptionSchedulePhaseItem { PriceId = "price_new", Quantity = 1 }],
+                                Discounts = [new SubscriptionSchedulePhaseDiscount { CouponId = "milestone-2c" }],
+                                ProrationBehavior = "none"
+                            }
+                        }
+                    }
+                ]
+            });
+
+        // Act
+        await _sut.HandleAsync(parsedEvent);
+
+        // Assert — schedule updated: Phase 1 skipped, Phase 2 included with cleared discounts
+        await _stripeAdapter.Received(1).UpdateSubscriptionScheduleAsync(
+            Arg.Is("sub_sched_789"),
+            Arg.Is<SubscriptionScheduleUpdateOptions>(o =>
+                o.DefaultSettings.AutomaticTax.Enabled == true &&
+                o.Phases.Count == 1 &&
+                o.Phases[0].AutomaticTax.Enabled == true &&
+                o.Phases[0].Items[0].Price == "price_new" &&
+                o.Phases[0].Discounts.Count == 0));
+    }
+
     #endregion
 }
 
diff --git a/test/Core.Test/Billing/Payment/Commands/UpdateBillingAddressCommandTests.cs b/test/Core.Test/Billing/Payment/Commands/UpdateBillingAddressCommandTests.cs
