fix(ui): prevent cloud upgrade modal flash on close (#12067)

ui/changelog.d/prowler-2254-cloud-upgrade-modal-flash.fixed.md

@@ -0,0 +1 @@
+Contextual Cloud upgrade modal content remains stable throughout the closing animation

ui/components/shared/cloud-upgrade-modal.test.tsx

@@ -7,9 +7,40 @@ import { CLOUD_UPGRADE_FEATURE } from "@/types/cloud-upgrade";
 
 import { CloudUpgradeModal } from "./cloud-upgrade-modal";
 
+const modalTestState = vi.hoisted(() => ({
+  keepContentMounted: false,
+}));
+
+vi.mock("@/components/shadcn/modal", async (importOriginal) => {
+  const actual =
+    await importOriginal<typeof import("@/components/shadcn/modal")>();
+  const { createElement } = await import("react");
+
+  return {
+    ...actual,
+    Modal: (props: Parameters<typeof actual.Modal>[0]) => {
+      if (!modalTestState.keepContentMounted) {
+        return createElement(actual.Modal, props);
+      }
+
+      return createElement(
+        "div",
+        { "aria-label": props.title, role: "dialog" },
+        createElement(
+          "button",
+          { onClick: () => props.onOpenChange?.(false), type: "button" },
+          "Close",
+        ),
+        props.children,
+      );
+    },
+  };
+});
+
 describe("CloudUpgradeModal", () => {
   afterEach(() => {
     cleanup();
+    modalTestState.keepContentMounted = false;
     vi.unstubAllEnvs();
     useCloudUpgradeStore.getState().closeCloudUpgrade();
   });
@@ -120,6 +151,42 @@ describe("CloudUpgradeModal", () => {
     expect(useCloudUpgradeStore.getState().activeFeature).toBeNull();
   });
 
+  it.each([
+    {
+      feature: CLOUD_UPGRADE_FEATURE.ALERTS,
+      otherTitle: "Add Your Entire AWS Organization",
+      title: "Turn Findings into Alerts",
+    },
+    {
+      feature: CLOUD_UPGRADE_FEATURE.AWS_ORGANIZATIONS,
+      otherTitle: "Turn Findings into Alerts",
+      title: "Add Your Entire AWS Organization",
+    },
+  ])(
+    "does not replace $title with another upgrade while closing",
+    async ({ feature, otherTitle, title }) => {
+      // Given
+      vi.stubEnv("NEXT_PUBLIC_IS_CLOUD_ENV", "false");
+      modalTestState.keepContentMounted = true;
+      const user = userEvent.setup();
+      useCloudUpgradeStore.getState().openCloudUpgrade(feature);
+
+      render(<CloudUpgradeModal />);
+      expect(screen.getByRole("dialog", { name: title })).toBeVisible();
+
+      // When
+      await user.click(screen.getByRole("button", { name: "Close" }));
+
+      // Then
+      expect(useCloudUpgradeStore.getState().activeFeature).toBeNull();
+      expect(screen.getByRole("dialog", { name: title })).toBeVisible();
+      expect(
+        screen.queryByText("Scale Prowler Without Operating It"),
+      ).not.toBeInTheDocument();
+      expect(screen.queryByText(otherTitle)).not.toBeInTheDocument();
+    },
+  );
+
   it("does not render upgrade UI in Prowler Cloud", () => {
     // Given
     vi.stubEnv("NEXT_PUBLIC_IS_CLOUD_ENV", "true");

ui/components/shared/cloud-upgrade-modal.tsx

@@ -14,12 +14,14 @@ import {
 } from "@/lib/cloud-upgrade";
 import { isCloud } from "@/lib/shared/env";
 import { useCloudUpgradeStore } from "@/store";
-import { CLOUD_UPGRADE_FEATURE } from "@/types/cloud-upgrade";
 
 const allowInitialAutoFocus = () => {};
 
 export const CloudUpgradeModal = () => {
   const activeFeature = useCloudUpgradeStore((state) => state.activeFeature);
+  const retainedFeature = useCloudUpgradeStore(
+    (state) => state.retainedFeature,
+  );
   const closeCloudUpgrade = useCloudUpgradeStore(
     (state) => state.closeCloudUpgrade,
   );
@@ -29,7 +31,7 @@ export const CloudUpgradeModal = () => {
 
   if (isCloud()) return null;
 
-  const feature = activeFeature ?? CLOUD_UPGRADE_FEATURE.GENERAL;
+  const feature = activeFeature ?? retainedFeature;
   const content = CLOUD_UPGRADE_CONTENT[feature];
 
   return (

ui/store/cloud-upgrade/store.test.ts

@@ -0,0 +1,51 @@
+import { beforeEach, describe, expect, it } from "vitest";
+
+import { CLOUD_UPGRADE_FEATURE } from "@/types/cloud-upgrade";
+
+import { useCloudUpgradeStore } from "./store";
+
+describe("useCloudUpgradeStore", () => {
+  beforeEach(() => {
+    useCloudUpgradeStore.setState({
+      activeFeature: null,
+      retainedFeature: CLOUD_UPGRADE_FEATURE.GENERAL,
+      returnFocusElement: null,
+    });
+  });
+
+  it("retains the opened feature when the modal closes", () => {
+    // Given
+    useCloudUpgradeStore
+      .getState()
+      .openCloudUpgrade(CLOUD_UPGRADE_FEATURE.ALERTS);
+
+    // When
+    useCloudUpgradeStore.getState().closeCloudUpgrade();
+
+    // Then
+    expect(useCloudUpgradeStore.getState().activeFeature).toBeNull();
+    expect(useCloudUpgradeStore.getState().retainedFeature).toBe(
+      CLOUD_UPGRADE_FEATURE.ALERTS,
+    );
+  });
+
+  it("updates the retained feature when another upgrade opens", () => {
+    // Given
+    useCloudUpgradeStore
+      .getState()
+      .openCloudUpgrade(CLOUD_UPGRADE_FEATURE.ALERTS);
+
+    // When
+    useCloudUpgradeStore
+      .getState()
+      .openCloudUpgrade(CLOUD_UPGRADE_FEATURE.AWS_ORGANIZATIONS);
+
+    // Then
+    expect(useCloudUpgradeStore.getState().activeFeature).toBe(
+      CLOUD_UPGRADE_FEATURE.AWS_ORGANIZATIONS,
+    );
+    expect(useCloudUpgradeStore.getState().retainedFeature).toBe(
+      CLOUD_UPGRADE_FEATURE.AWS_ORGANIZATIONS,
+    );
+  });
+});

ui/store/cloud-upgrade/store.ts

@@ -1,9 +1,13 @@
 import { create } from "zustand";
 
-import type { CloudUpgradeFeature } from "@/types/cloud-upgrade";
+import {
+  CLOUD_UPGRADE_FEATURE,
+  type CloudUpgradeFeature,
+} from "@/types/cloud-upgrade";
 
 interface CloudUpgradeStoreState {
   activeFeature: CloudUpgradeFeature | null;
+  retainedFeature: CloudUpgradeFeature;
   returnFocusElement: HTMLElement | null;
   openCloudUpgrade: (
     feature: CloudUpgradeFeature,
@@ -15,10 +19,12 @@ interface CloudUpgradeStoreState {
 // Upgrade prompts are ephemeral and shared so only one modal can be open.
 export const useCloudUpgradeStore = create<CloudUpgradeStoreState>((set) => ({
   activeFeature: null,
+  retainedFeature: CLOUD_UPGRADE_FEATURE.GENERAL,
   returnFocusElement: null,
   openCloudUpgrade: (activeFeature, requestedReturnFocusElement) =>
     set({
       activeFeature,
+      retainedFeature: activeFeature,
       returnFocusElement:
         requestedReturnFocusElement ??
         (document.activeElement instanceof HTMLElement