feat(biilling): Update unified upgrade dialog logic

apps/web/src/app/billing/individual/upgrade/unified-upgrade-dialog/unified-upgrade-dialog.component.html

@@ -1,15 +1,31 @@
 @if (step() == PlanSelectionStep) {
-  <app-upgrade-account
-    [dialogTitleMessageOverride]="planSelectionStepTitleOverride()"
-    [hideContinueWithoutUpgradingButton]="hideContinueWithoutUpgradingButton()"
-    (planSelected)="onPlanSelected($event)"
-    (closeClicked)="onCloseClicked()"
-  />
+  @if (hasPremiumPersonally()) {
+    <app-premium-org-upgrade
+      (planSelected)="onPlanSelected($event)"
+      (closeClicked)="onCloseClicked()"
+    />
+  } @else {
+    <app-upgrade-account
+      [dialogTitleMessageOverride]="planSelectionStepTitleOverride()"
+      [hideContinueWithoutUpgradingButton]="hideContinueWithoutUpgradingButton()"
+      (planSelected)="onPlanSelected($event)"
+      (closeClicked)="onCloseClicked()"
+    />
+  }
 } @else if (step() == PaymentStep && selectedPlan() !== null && account() !== null) {
-  <app-upgrade-payment
-    [selectedPlanId]="selectedPlan()"
-    [account]="account()"
-    (goBack)="previousStep()"
-    (complete)="onComplete($event)"
-  />
+  @if (hasPremiumPersonally()) {
+    <app-premium-org-upgrade-payment
+      [selectedPlanId]="selectedPlan()"
+      [account]="account()"
+      (goBack)="previousStep()"
+      (complete)="onComplete($event)"
+    />
+  } @else {
+    <app-upgrade-payment
+      [selectedPlanId]="selectedPersonalPlanId()"
+      [account]="account()"
+      (goBack)="previousStep()"
+      (complete)="onComplete($event)"
+    />
+  }
 }

apps/web/src/app/billing/individual/upgrade/unified-upgrade-dialog/unified-upgrade-dialog.component.spec.ts

@@ -3,17 +3,26 @@ import { ComponentFixture, TestBed } from "@angular/core/testing";
 import { NoopAnimationsModule } from "@angular/platform-browser/animations";
 import { Router } from "@angular/router";
 import { mock } from "jest-mock-extended";
+import { of } from "rxjs";
 
 import { PremiumInterestStateService } from "@bitwarden/angular/billing/services/premium-interest/premium-interest-state.service.abstraction";
 import { Account } from "@bitwarden/common/auth/abstractions/account.service";
+import { BillingAccountProfileStateService } from "@bitwarden/common/billing/abstractions";
 import {
   PersonalSubscriptionPricingTierId,
   PersonalSubscriptionPricingTierIds,
+  BusinessSubscriptionPricingTierId,
 } from "@bitwarden/common/billing/types/subscription-pricing-tier";
 import { mockAccountInfoWith } from "@bitwarden/common/spec";
 import { UserId } from "@bitwarden/common/types/guid";
 import { DIALOG_DATA, DialogRef } from "@bitwarden/components";
 
+import { PremiumOrgUpgradeComponent } from "../premium-org-upgrade/premium-org-upgrade.component";
+import {
+  PremiumOrgUpgradePaymentComponent,
+  PremiumOrgUpgradePaymentResult,
+  PremiumOrgUpgradePaymentStatus,
+} from "../premium-org-upgrade-payment/premium-org-upgrade-payment.component";
 import {
   UpgradeAccountComponent,
   UpgradeAccountStatus,
@@ -33,6 +42,7 @@ import {
   selector: "app-upgrade-account",
   template: "",
   standalone: true,
+  providers: [UpgradeAccountComponent],
   changeDetection: ChangeDetectionStrategy.OnPush,
 })
 class MockUpgradeAccountComponent {
@@ -46,6 +56,7 @@ class MockUpgradeAccountComponent {
   selector: "app-upgrade-payment",
   template: "",
   standalone: true,
+  providers: [UpgradePaymentComponent],
   changeDetection: ChangeDetectionStrategy.OnPush,
 })
 class MockUpgradePaymentComponent {
@@ -55,13 +66,43 @@ class MockUpgradePaymentComponent {
   complete = output<UpgradePaymentResult>();
 }
 
+@Component({
+  selector: "app-premium-org-upgrade",
+  template: "",
+  standalone: true,
+  providers: [PremiumOrgUpgradeComponent],
+  changeDetection: ChangeDetectionStrategy.OnPush,
+})
+class MockPremiumOrgUpgradeComponent {
+  readonly dialogTitleMessageOverride = input<string | null>(null);
+  readonly hideContinueWithoutUpgradingButton = input<boolean>(false);
+  planSelected = output<BusinessSubscriptionPricingTierId>();
+  closeClicked = output<PremiumOrgUpgradePaymentStatus>();
+}
+
+@Component({
+  selector: "app-premium-org-upgrade-payment",
+  template: "",
+  standalone: true,
+  providers: [PremiumOrgUpgradePaymentComponent],
+  changeDetection: ChangeDetectionStrategy.OnPush,
+})
+class MockPremiumOrgUpgradePaymentComponent {
+  readonly selectedPlanId = input<
+    PersonalSubscriptionPricingTierId | BusinessSubscriptionPricingTierId | null
+  >(null);
+  readonly account = input<Account | null>(null);
+  goBack = output<void>();
+  complete = output<{ status: PremiumOrgUpgradePaymentStatus; organizationId: string | null }>();
+}
+
 describe("UnifiedUpgradeDialogComponent", () => {
   let component: UnifiedUpgradeDialogComponent;
   let fixture: ComponentFixture<UnifiedUpgradeDialogComponent>;
   const mockDialogRef = mock<DialogRef>();
   const mockRouter = mock<Router>();
   const mockPremiumInterestStateService = mock<PremiumInterestStateService>();
-
+  const mockBillingAccountProfileStateService = mock<BillingAccountProfileStateService>();
   const mockAccount: Account = {
     id: "user-id" as UserId,
     ...mockAccountInfoWith({
@@ -97,14 +138,28 @@ describe("UnifiedUpgradeDialogComponent", () => {
         { provide: DIALOG_DATA, useValue: dialogData },
         { provide: Router, useValue: mockRouter },
         { provide: PremiumInterestStateService, useValue: mockPremiumInterestStateService },
+        {
+          provide: BillingAccountProfileStateService,
+          useValue: mockBillingAccountProfileStateService,
+        },
       ],
     })
       .overrideComponent(UnifiedUpgradeDialogComponent, {
         remove: {
-          imports: [UpgradeAccountComponent, UpgradePaymentComponent],
+          imports: [
+            UpgradeAccountComponent,
+            UpgradePaymentComponent,
+            PremiumOrgUpgradeComponent,
+            PremiumOrgUpgradePaymentComponent,
+          ],
         },
         add: {
-          imports: [MockUpgradeAccountComponent, MockUpgradePaymentComponent],
+          imports: [
+            MockUpgradeAccountComponent,
+            MockUpgradePaymentComponent,
+            MockPremiumOrgUpgradeComponent,
+            MockPremiumOrgUpgradePaymentComponent,
+          ],
         },
       })
       .compileComponents();
@@ -126,22 +181,36 @@ describe("UnifiedUpgradeDialogComponent", () => {
 
     // Default mock: no premium interest
     mockPremiumInterestStateService.getPremiumInterest.mockResolvedValue(false);
-
+    mockBillingAccountProfileStateService.hasPremiumPersonally$.mockReturnValue(of(true));
     await TestBed.configureTestingModule({
-      imports: [NoopAnimationsModule, UnifiedUpgradeDialogComponent],
+      imports: [UnifiedUpgradeDialogComponent],
       providers: [
         { provide: DialogRef, useValue: mockDialogRef },
         { provide: DIALOG_DATA, useValue: defaultDialogData },
         { provide: Router, useValue: mockRouter },
         { provide: PremiumInterestStateService, useValue: mockPremiumInterestStateService },
+        {
+          provide: BillingAccountProfileStateService,
+          useValue: mockBillingAccountProfileStateService,
+        },
       ],
     })
       .overrideComponent(UnifiedUpgradeDialogComponent, {
         remove: {
-          imports: [UpgradeAccountComponent, UpgradePaymentComponent],
+          imports: [
+            UpgradeAccountComponent,
+            UpgradePaymentComponent,
+            PremiumOrgUpgradeComponent,
+            PremiumOrgUpgradePaymentComponent,
+          ],
         },
         add: {
-          imports: [MockUpgradeAccountComponent, MockUpgradePaymentComponent],
+          imports: [
+            MockUpgradeAccountComponent,
+            MockUpgradePaymentComponent,
+            MockPremiumOrgUpgradeComponent,
+            MockPremiumOrgUpgradePaymentComponent,
+          ],
         },
       })
       .compileComponents();
@@ -401,4 +470,118 @@ describe("UnifiedUpgradeDialogComponent", () => {
       expect(mockDialogRef.close).toHaveBeenCalledWith({ status: "closed" });
     });
   });
+
+  describe("Premium Org Upgrade edge cases", () => {
+    it("should handle selecting a business plan (Teams) and move to payment step", async () => {
+      const { component } = await createComponentWithDialogData(defaultDialogData);
+
+      component["onPlanSelected"]("teams" as BusinessSubscriptionPricingTierId);
+
+      expect(component["selectedPlan"]()).toBe("teams");
+      expect(component["step"]()).toBe(UnifiedUpgradeDialogStep.Payment);
+    });
+
+    it("should handle completing premium org upgrade to Teams successfully", async () => {
+      const { component } = await createComponentWithDialogData(defaultDialogData);
+      mockRouter.navigate.mockResolvedValue(true);
+
+      const result = {
+        status: "upgradedToTeams" as const,
+        organizationId: "org-123",
+      };
+
+      await component["onComplete"](result);
+
+      expect(mockDialogRef.close).toHaveBeenCalledWith({
+        status: "upgradedToTeams",
+        organizationId: "org-123",
+      });
+    });
+
+    it("should handle completing premium org upgrade to Enterprise successfully", async () => {
+      const { component } = await createComponentWithDialogData(defaultDialogData);
+      mockRouter.navigate.mockResolvedValue(true);
+
+      const result = {
+        status: "upgradedToEnterprise" as const,
+        organizationId: "org-456",
+      };
+
+      await component["onComplete"](result);
+
+      expect(mockDialogRef.close).toHaveBeenCalledWith({
+        status: "upgradedToEnterprise",
+        organizationId: "org-456",
+      });
+    });
+
+    it("should handle user closing during premium org plan selection", async () => {
+      const { component } = await createComponentWithDialogData(defaultDialogData);
+
+      await component["onCloseClicked"]();
+
+      expect(mockDialogRef.close).toHaveBeenCalledWith({ status: "closed" });
+    });
+
+    it("should go back from premium org payment step to plan selection", async () => {
+      const { component } = await createComponentWithDialogData(defaultDialogData);
+      component["step"].set(UnifiedUpgradeDialogStep.Payment);
+      component["selectedPlan"].set("teams" as BusinessSubscriptionPricingTierId);
+
+      await component["previousStep"]();
+
+      expect(component["step"]()).toBe(UnifiedUpgradeDialogStep.PlanSelection);
+      expect(component["selectedPlan"]()).toBeNull();
+    });
+
+    it("should initialize with business plan when specified", async () => {
+      const customDialogData: UnifiedUpgradeDialogParams = {
+        account: mockAccount,
+        initialStep: UnifiedUpgradeDialogStep.Payment,
+        selectedPlan: "teams" as BusinessSubscriptionPricingTierId,
+      };
+
+      const { component: customComponent } = await createComponentWithDialogData(customDialogData);
+
+      expect(customComponent["step"]()).toBe(UnifiedUpgradeDialogStep.Payment);
+      expect(customComponent["selectedPlan"]()).toBe("teams");
+    });
+
+    it("should handle closed status during premium org upgrade", async () => {
+      const { component } = await createComponentWithDialogData(defaultDialogData);
+
+      const result: PremiumOrgUpgradePaymentResult = { status: "closed", organizationId: null };
+
+      await component["onComplete"](result);
+
+      expect(mockDialogRef.close).toHaveBeenCalledWith({
+        status: "closed",
+        organizationId: null,
+      });
+    });
+
+    it("should handle redirectOnCompletion for families upgrade with organization", async () => {
+      const customDialogData: UnifiedUpgradeDialogParams = {
+        account: mockAccount,
+        redirectOnCompletion: true,
+      };
+
+      mockRouter.navigate.mockResolvedValue(true);
+
+      const { component: customComponent } = await createComponentWithDialogData(customDialogData);
+
+      const result = {
+        status: "upgradedToFamilies" as const,
+        organizationId: "org-789",
+      };
+
+      await customComponent["onComplete"](result);
+
+      expect(mockRouter.navigate).toHaveBeenCalledWith(["/organizations/org-789/vault"]);
+      expect(mockDialogRef.close).toHaveBeenCalledWith({
+        status: "upgradedToFamilies",
+        organizationId: "org-789",
+      });
+    });
+  });
 });

apps/web/src/app/billing/individual/upgrade/unified-upgrade-dialog/unified-upgrade-dialog.component.ts

@@ -1,11 +1,23 @@
 import { DIALOG_DATA } from "@angular/cdk/dialog";
 import { CommonModule } from "@angular/common";
-import { ChangeDetectionStrategy, Component, Inject, OnInit, signal } from "@angular/core";
+import {
+  ChangeDetectionStrategy,
+  Component,
+  computed,
+  Inject,
+  OnInit,
+  signal,
+} from "@angular/core";
+import { toSignal } from "@angular/core/rxjs-interop";
 import { Router } from "@angular/router";
 
 import { PremiumInterestStateService } from "@bitwarden/angular/billing/services/premium-interest/premium-interest-state.service.abstraction";
 import { Account } from "@bitwarden/common/auth/abstractions/account.service";
-import { PersonalSubscriptionPricingTierId } from "@bitwarden/common/billing/types/subscription-pricing-tier";
+import { BillingAccountProfileStateService } from "@bitwarden/common/billing/abstractions";
+import {
+  BusinessSubscriptionPricingTierId,
+  PersonalSubscriptionPricingTierId,
+} from "@bitwarden/common/billing/types/subscription-pricing-tier";
 import { UnionOfValues } from "@bitwarden/common/vault/types/union-of-values";
 import {
   ButtonModule,
@@ -17,6 +29,11 @@ import {
 
 import { AccountBillingClient, PreviewInvoiceClient } from "../../../clients";
 import { BillingServicesModule } from "../../../services";
+import { PremiumOrgUpgradeComponent } from "../premium-org-upgrade/premium-org-upgrade.component";
+import {
+  PremiumOrgUpgradePaymentComponent,
+  PremiumOrgUpgradePaymentResult,
+} from "../premium-org-upgrade-payment/premium-org-upgrade-payment.component";
 import { UpgradeAccountComponent } from "../upgrade-account/upgrade-account.component";
 import { UpgradePaymentService } from "../upgrade-payment/services/upgrade-payment.service";
 import {
@@ -28,6 +45,8 @@ export const UnifiedUpgradeDialogStatus = {
   Closed: "closed",
   UpgradedToPremium: "upgradedToPremium",
   UpgradedToFamilies: "upgradedToFamilies",
+  UpgradedToTeams: "upgradedToTeams",
+  UpgradedToEnterprise: "upgradedToEnterprise",
 } as const;
 
 export const UnifiedUpgradeDialogStep = {
@@ -57,7 +76,7 @@ export type UnifiedUpgradeDialogResult = {
 export type UnifiedUpgradeDialogParams = {
   account: Account;
   initialStep?: UnifiedUpgradeDialogStep | null;
-  selectedPlan?: PersonalSubscriptionPricingTierId | null;
+  selectedPlan?: PersonalSubscriptionPricingTierId | BusinessSubscriptionPricingTierId | null;
   planSelectionStepTitleOverride?: string | null;
   hideContinueWithoutUpgradingButton?: boolean;
   redirectOnCompletion?: boolean;
@@ -73,6 +92,8 @@ export type UnifiedUpgradeDialogParams = {
     UpgradeAccountComponent,
     UpgradePaymentComponent,
     BillingServicesModule,
+    PremiumOrgUpgradeComponent,
+    PremiumOrgUpgradePaymentComponent,
   ],
   providers: [UpgradePaymentService, AccountBillingClient, PreviewInvoiceClient],
   templateUrl: "./unified-upgrade-dialog.component.html",
@@ -82,11 +103,23 @@ export class UnifiedUpgradeDialogComponent implements OnInit {
   protected readonly step = signal<UnifiedUpgradeDialogStep>(
     UnifiedUpgradeDialogStep.PlanSelection,
   );
-  protected readonly selectedPlan = signal<PersonalSubscriptionPricingTierId | null>(null);
+  protected readonly selectedPlan = signal<
+    PersonalSubscriptionPricingTierId | BusinessSubscriptionPricingTierId | null
+  >(null);
   protected readonly account = signal<Account | null>(null);
   protected readonly planSelectionStepTitleOverride = signal<string | null>(null);
   protected readonly hideContinueWithoutUpgradingButton = signal<boolean>(false);
   protected readonly hasPremiumInterest = signal(false);
+  protected readonly hasPremiumPersonally = toSignal(
+    this.billingAccountProfileStateService.hasPremiumPersonally$(this.params.account.id),
+    { initialValue: false },
+  );
+
+  // Type-narrowed computed signal for app-upgrade-payment
+  // When hasPremiumPersonally is false, selectedPlan will only contain PersonalSubscriptionPricingTierId
+  protected readonly selectedPersonalPlanId = computed<PersonalSubscriptionPricingTierId | null>(
+    () => this.selectedPlan() as PersonalSubscriptionPricingTierId | null,
+  );
 
   protected readonly PaymentStep = UnifiedUpgradeDialogStep.Payment;
   protected readonly PlanSelectionStep = UnifiedUpgradeDialogStep.PlanSelection;
@@ -96,6 +129,7 @@ export class UnifiedUpgradeDialogComponent implements OnInit {
     @Inject(DIALOG_DATA) private params: UnifiedUpgradeDialogParams,
     private router: Router,
     private premiumInterestStateService: PremiumInterestStateService,
+    private billingAccountProfileStateService: BillingAccountProfileStateService,
   ) {}
 
   async ngOnInit(): Promise<void> {
@@ -121,7 +155,9 @@ export class UnifiedUpgradeDialogComponent implements OnInit {
     }
   }
 
-  protected onPlanSelected(planId: PersonalSubscriptionPricingTierId): void {
+  protected onPlanSelected(
+    planId: PersonalSubscriptionPricingTierId | BusinessSubscriptionPricingTierId,
+  ): void {
     this.selectedPlan.set(planId);
     this.nextStep();
   }
@@ -150,9 +186,17 @@ export class UnifiedUpgradeDialogComponent implements OnInit {
     }
   }
 
-  protected async onComplete(result: UpgradePaymentResult): Promise<void> {
+  protected async onComplete(
+    result: UpgradePaymentResult | PremiumOrgUpgradePaymentResult,
+  ): Promise<void> {
     let status: UnifiedUpgradeDialogStatus;
     switch (result.status) {
+      case "upgradedToTeams":
+        status = UnifiedUpgradeDialogStatus.UpgradedToTeams;
+        break;
+      case "upgradedToEnterprise":
+        status = UnifiedUpgradeDialogStatus.UpgradedToEnterprise;
+        break;
       case "upgradedToPremium":
         status = UnifiedUpgradeDialogStatus.UpgradedToPremium;
         break;