feat(react,shadcn): Add useMultiFactorAssertionCleanup utility hook

Author: Ehesp

packages/react/src/auth/forms/multi-factor-auth-assertion-form.test.tsx

@@ -14,9 +14,12 @@
  */
 
 import { describe, it, expect, vi, afterEach } from "vitest";
-import { render, screen, fireEvent, cleanup } from "@testing-library/react";
-import { MultiFactorAuthAssertionForm } from "~/auth/forms/multi-factor-auth-assertion-form";
-import { CreateFirebaseUIProvider, createMockUI } from "~/tests/utils";
+import { render, screen, fireEvent, cleanup, renderHook } from "@testing-library/react";
+import {
+  MultiFactorAuthAssertionForm,
+  useMultiFactorAssertionCleanup,
+} from "~/auth/forms/multi-factor-auth-assertion-form";
+import { CreateFirebaseUIProvider, createMockUI, createFirebaseUIProvider } from "~/tests/utils";
 import { registerLocale } from "@invertase/firebaseui-translations";
 import { FactorId, MultiFactorResolver, PhoneMultiFactorGenerator, TotpMultiFactorGenerator } from "firebase/auth";
 
@@ -55,6 +58,46 @@ afterEach(() => {
   vi.clearAllMocks();
 });
 
+describe("useMultiFactorAssertionCleanup", () => {
+  it("calls setMultiFactorResolver on unmount", () => {
+    const ui = createMockUI();
+    const setMultiFactorResolverSpy = vi.spyOn(ui.get(), "setMultiFactorResolver");
+
+    const { unmount } = renderHook(() => useMultiFactorAssertionCleanup(), {
+      wrapper: ({ children }) => createFirebaseUIProvider({ children, ui }),
+    });
+
+    expect(setMultiFactorResolverSpy).not.toHaveBeenCalled();
+
+    unmount();
+
+    expect(setMultiFactorResolverSpy).toHaveBeenCalledTimes(1);
+  });
+
+  it("clears multiFactorResolver when component unmounts", () => {
+    const ui = createMockUI();
+    const mockResolver = {
+      auth: {} as any,
+      session: null,
+      hints: [],
+    } as unknown as MultiFactorResolver;
+    ui.get().setMultiFactorResolver(mockResolver);
+
+    const setMultiFactorResolverSpy = vi.spyOn(ui.get(), "setMultiFactorResolver");
+
+    const { unmount } = renderHook(() => useMultiFactorAssertionCleanup(), {
+      wrapper: ({ children }) => createFirebaseUIProvider({ children, ui }),
+    });
+
+    expect(ui.get().multiFactorResolver).toBe(mockResolver);
+
+    unmount();
+
+    expect(setMultiFactorResolverSpy).toHaveBeenCalledTimes(1);
+    expect(ui.get().multiFactorResolver).toBeUndefined();
+  });
+});
+
 describe("<MultiFactorAuthAssertionForm />", () => {
   it("throws error when no multiFactorResolver is present", () => {
     const ui = createMockUI();

packages/react/src/auth/forms/multi-factor-auth-assertion-form.tsx

@@ -4,7 +4,7 @@ import {
   type UserCredential,
   type MultiFactorInfo,
 } from "firebase/auth";
-import { type ComponentProps, useState } from "react";
+import { type ComponentProps, useEffect, useState } from "react";
 import { useUI } from "~/hooks";
 import { TotpMultiFactorAssertionForm } from "../forms/mfa/totp-multi-factor-assertion-form";
 import { SmsMultiFactorAssertionForm } from "../forms/mfa/sms-multi-factor-assertion-form";
@@ -15,11 +15,24 @@ export type MultiFactorAuthAssertionFormProps = {
   onSuccess?: (credential: UserCredential) => void;
 };
 
+export function useMultiFactorAssertionCleanup() {
+  const ui = useUI();
+
+  useEffect(() => {
+    return () => {
+      ui.setMultiFactorResolver();
+    };
+    // eslint-disable-next-line react-hooks/exhaustive-deps -- UI isn't stable enough to be a dependency here. Could we use useEffectEvent here instead once we depend on 19.2?
+  }, []);
+}
+
 export function MultiFactorAuthAssertionForm(props: MultiFactorAuthAssertionFormProps) {
   const ui = useUI();
   const resolver = ui.multiFactorResolver;
   const mfaAssertionFactorPrompt = getTranslation(ui, "prompts", "mfaAssertionFactorPrompt");
 
+  useMultiFactorAssertionCleanup();
+
   if (!resolver) {
     throw new Error("MultiFactorAuthAssertionForm requires a multi-factor resolver");
   }

packages/react/src/auth/index.ts

@@ -33,6 +33,7 @@ export {
 } from "./forms/forgot-password-auth-form";
 export {
   MultiFactorAuthAssertionForm,
+  useMultiFactorAssertionCleanup,
   type MultiFactorAuthAssertionFormProps,
 } from "./forms/multi-factor-auth-assertion-form";
 export {

packages/shadcn/src/components/multi-factor-auth-assertion-form.test.tsx

@@ -3,7 +3,16 @@ import { render, screen, fireEvent, cleanup } from "@testing-library/react";
 import { MultiFactorAuthAssertionForm } from "./multi-factor-auth-assertion-form";
 import { createFirebaseUIProvider, createMockUI } from "../../tests/utils";
 import { registerLocale } from "@invertase/firebaseui-translations";
-import { FactorId, MultiFactorResolver, PhoneMultiFactorGenerator, TotpMultiFactorGenerator } from "firebase/auth";
+import { MultiFactorResolver, PhoneMultiFactorGenerator, TotpMultiFactorGenerator } from "firebase/auth";
+
+const mockUseMultiFactorAssertionCleanup = vi.fn();
+vi.mock("@invertase/firebaseui-react", async () => {
+  const actual = await vi.importActual<typeof import("@invertase/firebaseui-react")>("@invertase/firebaseui-react");
+  return {
+    ...actual,
+    useMultiFactorAssertionCleanup: () => mockUseMultiFactorAssertionCleanup(),
+  };
+});
 
 vi.mock("@/components/sms-multi-factor-assertion-form", () => ({
   SmsMultiFactorAssertionForm: ({ hint, onSuccess }: { hint: any; onSuccess?: (credential: any) => void }) => (
@@ -30,12 +39,37 @@ vi.mock("@/components/totp-multi-factor-assertion-form", () => ({
 describe("<MultiFactorAuthAssertionForm />", () => {
   beforeEach(() => {
     vi.clearAllMocks();
+    mockUseMultiFactorAssertionCleanup.mockClear();
   });
 
   afterEach(() => {
     cleanup();
   });
 
+  it("calls useMultiFactorAssertionCleanup when component renders", () => {
+    const mockResolver: MultiFactorResolver = {
+      hints: [
+        {
+          uid: "test-uid",
+          factorId: PhoneMultiFactorGenerator.FACTOR_ID,
+          displayName: "Test Phone",
+        },
+      ],
+    } as MultiFactorResolver;
+
+    const ui = createMockUI();
+    ui.get().setMultiFactorResolver(mockResolver as unknown as MultiFactorResolver);
+
+    render(
+      createFirebaseUIProvider({
+        children: <MultiFactorAuthAssertionForm />,
+        ui: ui,
+      })
+    );
+
+    expect(mockUseMultiFactorAssertionCleanup).toHaveBeenCalledTimes(1);
+  });
+
   it("throws error when no multiFactorResolver is present", () => {
     const ui = createMockUI();
 

packages/shadcn/src/components/multi-factor-auth-assertion-form.tsx

@@ -8,7 +8,8 @@ import {
   type MultiFactorInfo,
   type UserCredential,
 } from "firebase/auth";
-import { useEffect, useState, type ComponentProps } from "react";
+import { useState, type ComponentProps } from "react";
+import { useMultiFactorAssertionCleanup } from "@invertase/firebaseui-react";
 
 import { SmsMultiFactorAssertionForm } from "@/components/sms-multi-factor-assertion-form";
 import { TotpMultiFactorAssertionForm } from "@/components/totp-multi-factor-assertion-form";
@@ -23,11 +24,7 @@ export function MultiFactorAuthAssertionForm({ onSuccess }: MultiFactorAuthAsser
   const resolver = ui.multiFactorResolver;
   const mfaAssertionFactorPrompt = getTranslation(ui, "prompts", "mfaAssertionFactorPrompt");
 
-  useEffect(() => {
-    return () => {
-      ui.setMultiFactorResolver();
-    };
-  }, []); // eslint-disable-line react-hooks/exhaustive-deps
+  useMultiFactorAssertionCleanup();
 
   if (!resolver) {
     throw new Error("MultiFactorAuthAssertionForm requires a multi-factor resolver");