prowler/ui/components/providers/wizard/hooks/use-provider-wizard-controller.test.tsx

import { act, renderHook, waitFor } from "@testing-library/react";
import { beforeEach, describe, expect, it, vi } from "vitest";

import { useOrgSetupStore } from "@/store/organizations/store";
import { useProviderWizardStore } from "@/store/provider-wizard/store";
import { ORG_WIZARD_STEP } from "@/types/organizations";
import {
  PROVIDER_WIZARD_MODE,
  PROVIDER_WIZARD_STEP,
} from "@/types/provider-wizard";

import type { ProviderWizardInitialData } from "../types";
import { useProviderWizardController } from "./use-provider-wizard-controller";

const { refreshMock, requestOpenOnWizardCloseMock } = vi.hoisted(() => ({
  refreshMock: vi.fn(),
  requestOpenOnWizardCloseMock: vi.fn(),
}));

vi.mock("next/navigation", () => ({
  useRouter: () => ({
    refresh: refreshMock,
  }),
}));

vi.mock("@/store/onboarding-checkpoint", () => ({
  useOnboardingCheckpointStore: {
    getState: () => ({
      requestOpenOnWizardClose: requestOpenOnWizardCloseMock,
    }),
  },
}));

vi.mock("next-auth/react", () => ({
  useSession: () => ({
    data: null,
    status: "unauthenticated",
  }),
}));

describe("useProviderWizardController", () => {
  beforeEach(() => {
    vi.useRealTimers();
    vi.clearAllMocks();
    requestOpenOnWizardCloseMock.mockClear();
    sessionStorage.clear();
    localStorage.clear();
    // Checkpoint is Cloud-only.
    vi.stubEnv("NEXT_PUBLIC_IS_CLOUD_ENV", "true");
    useProviderWizardStore.getState().reset();
    useOrgSetupStore.getState().reset();
  });

  it("refreshes providers data when the wizard closes", () => {
    // Given
    const onOpenChange = vi.fn();
    const { result } = renderHook(() =>
      useProviderWizardController({
        open: true,
        onOpenChange,
      }),
    );

    // When
    act(() => {
      result.current.handleClose();
    });

    // Then
    expect(onOpenChange).toHaveBeenCalledWith(false);
    expect(refreshMock).toHaveBeenCalledTimes(1);
  });

  it("requests the onboarding checkpoint with providerConnected:true when a provider was created", () => {
    const onOpenChange = vi.fn();
    const { result } = renderHook(() =>
      useProviderWizardController({
        open: true,
        onOpenChange,
      }),
    );
    act(() => {
      useProviderWizardStore.getState().setProvider({
        id: "provider-1",
        type: "aws",
        uid: "111111111111",
        alias: "production",
      });
    });

    act(() => {
      result.current.handleClose();
    });

    // providerId is read before the store reset clears it.
    expect(requestOpenOnWizardCloseMock).toHaveBeenCalledWith({
      providerConnected: true,
    });
  });

  it("requests the onboarding checkpoint with providerConnected:false when no provider was created", () => {
    const onOpenChange = vi.fn();
    const { result } = renderHook(() =>
      useProviderWizardController({
        open: true,
        onOpenChange,
      }),
    );

    act(() => {
      result.current.handleClose();
    });

    expect(requestOpenOnWizardCloseMock).toHaveBeenCalledWith({
      providerConnected: false,
    });
  });

  it("does not request the onboarding checkpoint in self-hosted (OSS) deployments", () => {
    vi.stubEnv("NEXT_PUBLIC_IS_CLOUD_ENV", "false");
    const onOpenChange = vi.fn();
    const { result } = renderHook(() =>
      useProviderWizardController({
        open: true,
        onOpenChange,
      }),
    );
    act(() => {
      useProviderWizardStore.getState().setProvider({
        id: "provider-1",
        type: "aws",
        uid: "111111111111",
        alias: "production",
      });
    });

    act(() => {
      result.current.handleClose();
    });

    // Checkpoint stays untouched, but the close still refreshes.
    expect(requestOpenOnWizardCloseMock).not.toHaveBeenCalled();
    expect(refreshMock).toHaveBeenCalledTimes(1);
  });

  it("hydrates update mode when initial data is provided", async () => {
    // Given
    const onOpenChange = vi.fn();

    // When
    const { result } = renderHook(() =>
      useProviderWizardController({
        open: true,
        onOpenChange,
        initialData: {
          providerId: "provider-1",
          providerType: "aws",
          providerUid: "111111111111",
          providerAlias: "production",
          secretId: "secret-1",
          mode: PROVIDER_WIZARD_MODE.UPDATE,
        },
      }),
    );

    // Then
    await waitFor(() => {
      expect(result.current.currentStep).toBe(PROVIDER_WIZARD_STEP.CREDENTIALS);
    });
    expect(result.current.modalTitle).toBe("Update Provider Credentials");
    expect(result.current.isProviderFlow).toBe(true);
    expect(result.current.docsLink).toBe(
      "https://goto.prowler.com/provider-aws",
    );

    const state = useProviderWizardStore.getState();
    expect(state.providerId).toBe("provider-1");
    expect(state.providerType).toBe("aws");
    expect(state.providerUid).toBe("111111111111");
    expect(state.providerAlias).toBe("production");
    expect(state.secretId).toBe("secret-1");
    expect(state.mode).toBe(PROVIDER_WIZARD_MODE.UPDATE);
  });

  it("switches into and out of organizations flow", () => {
    // Given
    const onOpenChange = vi.fn();
    const { result } = renderHook(() =>
      useProviderWizardController({
        open: true,
        onOpenChange,
      }),
    );

    // When
    act(() => {
      result.current.openOrganizationsFlow();
    });

    // Then
    expect(result.current.wizardVariant).toBe("organizations");
    expect(result.current.isProviderFlow).toBe(false);
    expect(result.current.orgCurrentStep).toBe(ORG_WIZARD_STEP.SETUP);
    expect(result.current.docsLink).toBe(
      "https://docs.prowler.com/user-guide/tutorials/prowler-cloud-aws-organizations",
    );

    // When
    act(() => {
      result.current.backToProviderFlow();
    });

    // Then
    expect(result.current.wizardVariant).toBe("provider");
    expect(result.current.isProviderFlow).toBe(true);
    expect(result.current.currentStep).toBe(PROVIDER_WIZARD_STEP.CONNECT);
  });

  it("moves to launch step after a successful connection test in add mode", () => {
    // Given
    const onOpenChange = vi.fn();
    const { result } = renderHook(() =>
      useProviderWizardController({
        open: true,
        onOpenChange,
      }),
    );

    // When
    act(() => {
      result.current.setCurrentStep(PROVIDER_WIZARD_STEP.TEST);
      result.current.handleTestSuccess();
    });

    // Then
    expect(result.current.currentStep).toBe(PROVIDER_WIZARD_STEP.LAUNCH);
    expect(onOpenChange).not.toHaveBeenCalled();
  });

  it("closes the wizard after a successful connection test in update mode", async () => {
    // Given
    const onOpenChange = vi.fn();
    const { result } = renderHook(() =>
      useProviderWizardController({
        open: true,
        onOpenChange,
        initialData: {
          providerId: "provider-1",
          providerType: "aws",
          providerUid: "111111111111",
          providerAlias: "production",
          secretId: "secret-1",
          mode: PROVIDER_WIZARD_MODE.UPDATE,
        },
      }),
    );

    await waitFor(() => {
      expect(result.current.currentStep).toBe(PROVIDER_WIZARD_STEP.CREDENTIALS);
    });

    // When
    act(() => {
      result.current.setCurrentStep(PROVIDER_WIZARD_STEP.TEST);
      result.current.handleTestSuccess();
    });

    // Credential rotation skips the launch/schedule step.
    expect(onOpenChange).toHaveBeenCalledWith(false);
    expect(refreshMock).toHaveBeenCalledTimes(1);
    expect(result.current.currentStep).not.toBe(PROVIDER_WIZARD_STEP.LAUNCH);
  });

  it("does not override launch footer config in the controller", () => {
    // Given
    const onOpenChange = vi.fn();
    const { result } = renderHook(() =>
      useProviderWizardController({
        open: true,
        onOpenChange,
      }),
    );

    // When
    act(() => {
      result.current.setCurrentStep(PROVIDER_WIZARD_STEP.LAUNCH);
    });

    // Then
    expect(result.current.resolvedFooterConfig.showAction).toBe(false);
    expect(result.current.resolvedFooterConfig.showBack).toBe(false);
    expect(onOpenChange).not.toHaveBeenCalled();
  });

  it("does not reset organizations step when org store updates while modal is open", () => {
    // Given
    const onOpenChange = vi.fn();
    const { result } = renderHook(() =>
      useProviderWizardController({
        open: true,
        onOpenChange,
      }),
    );

    act(() => {
      result.current.openOrganizationsFlow();
      result.current.setOrgCurrentStep(ORG_WIZARD_STEP.VALIDATE);
    });

    // When
    act(() => {
      useOrgSetupStore
        .getState()
        .setOrganization("org-1", "My Org", "o-abc123def4");
      useOrgSetupStore.getState().setDiscovery("disc-1", {
        roots: [],
        organizational_units: [],
        accounts: [],
      });
    });

    // Then
    expect(result.current.wizardVariant).toBe("organizations");
    expect(result.current.orgCurrentStep).toBe(ORG_WIZARD_STEP.VALIDATE);
  });

  it("does not rehydrate wizard state when initial data changes while modal remains open", async () => {
    // Given
    const onOpenChange = vi.fn();
    const { result, rerender } = renderHook(
      ({
        open,
        initialData,
      }: {
        open: boolean;
        initialData?: ProviderWizardInitialData;
      }) =>
        useProviderWizardController({
          open,
          onOpenChange,
          initialData,
        }),
      {
        initialProps: {
          open: true,
          initialData: {
            providerId: "provider-1",
            providerType: "gcp",
            providerUid: "project-123",
            providerAlias: "gcp-main",
            secretId: null,
            mode: PROVIDER_WIZARD_MODE.ADD,
          },
        },
      },
    );

    await waitFor(() => {
      expect(result.current.currentStep).toBe(PROVIDER_WIZARD_STEP.CREDENTIALS);
    });

    act(() => {
      useProviderWizardStore.getState().setVia("service-account");
      result.current.setCurrentStep(PROVIDER_WIZARD_STEP.TEST);
    });

    // Provider data refreshes while modal is still open — user progress must be kept.
    rerender({
      open: true,
      initialData: {
        providerId: "provider-1",
        providerType: "gcp",
        providerUid: "project-123",
        providerAlias: "gcp-main",
        secretId: "secret-1",
        mode: PROVIDER_WIZARD_MODE.UPDATE,
      },
    });

    expect(result.current.currentStep).toBe(PROVIDER_WIZARD_STEP.TEST);
    expect(useProviderWizardStore.getState().via).toBe("service-account");
  });
});