Merge pull request #1211 from firebase/@invertase/anon-on-upgrade

Ehesp · web-flow · commit 1c78f8e0015e · 2025-10-02T15:25:56.000+01:00
diff --git a/packages/core/src/behaviors/anonymous-upgrade.test.ts b/packages/core/src/behaviors/anonymous-upgrade.test.ts
@@ -1,13 +1,18 @@
 import { describe, it, expect, vi, beforeEach } from "vitest";
-import { Auth, AuthCredential, AuthProvider, linkWithCredential, linkWithRedirect, User } from "firebase/auth";
-import { autoUpgradeAnonymousCredentialHandler, autoUpgradeAnonymousProviderHandler } from "./anonymous-upgrade";
+import { Auth, AuthCredential, AuthProvider, linkWithCredential, linkWithRedirect, User, UserCredential } from "firebase/auth";
+import { autoUpgradeAnonymousCredentialHandler, autoUpgradeAnonymousProviderHandler, autoUpgradeAnonymousUserRedirectHandler, OnUpgradeCallback } from "./anonymous-upgrade";
 import { createMockUI } from "~/tests/utils";
+import { getBehavior } from "~/behaviors";
 
 vi.mock("firebase/auth", () => ({
   linkWithCredential: vi.fn(),
   linkWithRedirect: vi.fn(),
 }));
 
+vi.mock("~/behaviors", () => ({
+  getBehavior: vi.fn(),
+}));
+
 beforeEach(() => {
   vi.clearAllMocks();
 });
@@ -30,6 +35,38 @@ describe("autoUpgradeAnonymousCredentialHandler", () => {
     expect(result).toBe(mockResult);
   });
 
+  it("should call onUpgrade callback when provided", async () => {
+    const mockUser = { isAnonymous: true, uid: "anonymous-123" } as User;
+    const mockAuth = { currentUser: mockUser } as Auth;
+    const mockUI = createMockUI({ auth: mockAuth });
+    const mockCredential = { providerId: "password" } as AuthCredential;
+    const mockResult = { user: { uid: "upgraded-123" } } as UserCredential;
+    
+    vi.mocked(linkWithCredential).mockResolvedValue(mockResult);
+    
+    const onUpgrade = vi.fn().mockResolvedValue(undefined);
+    
+    const result = await autoUpgradeAnonymousCredentialHandler(mockUI, mockCredential, onUpgrade);
+
+    expect(onUpgrade).toHaveBeenCalledWith(mockUI, "anonymous-123", mockResult);
+    expect(result).toBe(mockResult);
+  });
+
+  it("should handle onUpgrade callback errors", async () => {
+    const mockUser = { isAnonymous: true, uid: "anonymous-123" } as User;
+    const mockAuth = { currentUser: mockUser } as Auth;
+    const mockUI = createMockUI({ auth: mockAuth });
+    const mockCredential = { providerId: "password" } as AuthCredential;
+    const mockResult = { user: { uid: "upgraded-123" } } as UserCredential;
+    
+    vi.mocked(linkWithCredential).mockResolvedValue(mockResult);
+    
+    const onUpgrade = vi.fn().mockRejectedValue(new Error("Callback error"));
+    
+    await expect(autoUpgradeAnonymousCredentialHandler(mockUI, mockCredential, onUpgrade))
+      .rejects.toThrow("Callback error");
+  });
+
   it("should not upgrade when user is not anonymous", async () => {
     const mockUser = { isAnonymous: false, uid: "regular-user-123" } as User;
     const mockAuth = { currentUser: mockUser } as Auth;
@@ -62,14 +99,55 @@ describe("autoUpgradeAnonymousProviderHandler", () => {
     const mockAuth = { currentUser: mockUser } as Auth;
     const mockUI = createMockUI({ auth: mockAuth });
     const mockProvider = { providerId: "google.com" } as AuthProvider;
+    const mockResult = { user: { uid: "upgraded-123" } } as UserCredential;
 
-    vi.mocked(linkWithRedirect).mockResolvedValue({} as never);
+    const mockProviderLinkStrategy = vi.fn().mockResolvedValue(mockResult);
+    vi.mocked(getBehavior).mockReturnValue(mockProviderLinkStrategy);
 
-    await autoUpgradeAnonymousProviderHandler(mockUI, mockProvider);
+    const localStorageSpy = vi.spyOn(Storage.prototype, 'setItem');
+    const localStorageRemoveSpy = vi.spyOn(Storage.prototype, 'removeItem');
 
-    expect(linkWithRedirect).toHaveBeenCalledWith(mockUser, mockProvider);
-    expect(mockUI.setState).toHaveBeenCalledWith("pending");
-    expect(mockUI.setState).not.toHaveBeenCalledWith("idle");
+    const result = await autoUpgradeAnonymousProviderHandler(mockUI, mockProvider);
+
+    expect(getBehavior).toHaveBeenCalledWith(mockUI, "providerLinkStrategy");
+    expect(mockProviderLinkStrategy).toHaveBeenCalledWith(mockUI, mockUser, mockProvider);
+    expect(localStorageSpy).toHaveBeenCalledWith("fbui:upgrade:oldUserId", "anonymous-123");
+    expect(localStorageRemoveSpy).toHaveBeenCalledWith("fbui:upgrade:oldUserId");
+    expect(result).toBe(mockResult);
+  });
+
+  it("should call onUpgrade callback when provided", async () => {
+    const mockUser = { isAnonymous: true, uid: "anonymous-123" } as User;
+    const mockAuth = { currentUser: mockUser } as Auth;
+    const mockUI = createMockUI({ auth: mockAuth });
+    const mockProvider = { providerId: "google.com" } as AuthProvider;
+    const mockResult = { user: { uid: "upgraded-123" } } as UserCredential;
+
+    const mockProviderLinkStrategy = vi.fn().mockResolvedValue(mockResult);
+    vi.mocked(getBehavior).mockReturnValue(mockProviderLinkStrategy);
+
+    const onUpgrade = vi.fn().mockResolvedValue(undefined);
+
+    const result = await autoUpgradeAnonymousProviderHandler(mockUI, mockProvider, onUpgrade);
+
+    expect(onUpgrade).toHaveBeenCalledWith(mockUI, "anonymous-123", mockResult);
+    expect(result).toBe(mockResult);
+  });
+
+  it("should handle onUpgrade callback errors", async () => {
+    const mockUser = { isAnonymous: true, uid: "anonymous-123" } as User;
+    const mockAuth = { currentUser: mockUser } as Auth;
+    const mockUI = createMockUI({ auth: mockAuth });
+    const mockProvider = { providerId: "google.com" } as AuthProvider;
+    const mockResult = { user: { uid: "upgraded-123" } } as UserCredential;
+
+    const mockProviderLinkStrategy = vi.fn().mockResolvedValue(mockResult);
+    vi.mocked(getBehavior).mockReturnValue(mockProviderLinkStrategy);
+
+    const onUpgrade = vi.fn().mockRejectedValue(new Error("Callback error"));
+
+    await expect(autoUpgradeAnonymousProviderHandler(mockUI, mockProvider, onUpgrade))
+      .rejects.toThrow("Callback error");
   });
 
   it("should not upgrade when user is not anonymous", async () => {
@@ -95,3 +173,77 @@ describe("autoUpgradeAnonymousProviderHandler", () => {
     expect(mockUI.setState).not.toHaveBeenCalled();
   });
 });
+
+describe("autoUpgradeAnonymousUserRedirectHandler", () => {
+  beforeEach(() => {
+    window.localStorage.clear();
+  });
+
+  it("should call onUpgrade callback when oldUserId exists in localStorage", async () => {
+    const mockUI = createMockUI();
+    const mockCredential = { user: { uid: "upgraded-123" } } as UserCredential;
+    const oldUserId = "anonymous-123";
+    
+    window.localStorage.setItem("fbui:upgrade:oldUserId", oldUserId);
+    
+    const onUpgrade = vi.fn().mockResolvedValue(undefined);
+    
+    await autoUpgradeAnonymousUserRedirectHandler(mockUI, mockCredential, onUpgrade);
+    
+    expect(onUpgrade).toHaveBeenCalledWith(mockUI, oldUserId, mockCredential);
+    expect(window.localStorage.getItem("fbui:upgrade:oldUserId")).toBeNull();
+  });
+
+  it("should not call onUpgrade callback when no oldUserId in localStorage", async () => {
+    const mockUI = createMockUI();
+    const mockCredential = { user: { uid: "upgraded-123" } } as UserCredential;
+    
+    const onUpgrade = vi.fn().mockResolvedValue(undefined);
+    
+    await autoUpgradeAnonymousUserRedirectHandler(mockUI, mockCredential, onUpgrade);
+    
+    expect(onUpgrade).not.toHaveBeenCalled();
+  });
+
+  it("should not call onUpgrade callback when no credential provided", async () => {
+    const mockUI = createMockUI();
+    const oldUserId = "anonymous-123";
+    
+    window.localStorage.setItem("fbui:upgrade:oldUserId", oldUserId);
+    
+    const onUpgrade = vi.fn().mockResolvedValue(undefined);
+    
+    await autoUpgradeAnonymousUserRedirectHandler(mockUI, null, onUpgrade);
+    
+    expect(onUpgrade).not.toHaveBeenCalled();
+  });
+
+  it("should not call onUpgrade callback when no onUpgrade callback provided", async () => {
+    const mockUI = createMockUI();
+    const mockCredential = { user: { uid: "upgraded-123" } } as UserCredential;
+    const oldUserId = "anonymous-123";
+    
+    window.localStorage.setItem("fbui:upgrade:oldUserId", oldUserId);
+    
+    await autoUpgradeAnonymousUserRedirectHandler(mockUI, mockCredential);
+    
+    // Should not throw and should clean up localStorage even when no callback provided
+    expect(window.localStorage.getItem("fbui:upgrade:oldUserId")).toBeNull();
+  });
+
+  it("should handle onUpgrade callback errors", async () => {
+    const mockUI = createMockUI();
+    const mockCredential = { user: { uid: "upgraded-123" } } as UserCredential;
+    const oldUserId = "anonymous-123";
+    
+    window.localStorage.setItem("fbui:upgrade:oldUserId", oldUserId);
+    
+    const onUpgrade = vi.fn().mockRejectedValue(new Error("Callback error"));
+    
+    await expect(autoUpgradeAnonymousUserRedirectHandler(mockUI, mockCredential, onUpgrade))
+      .rejects.toThrow("Callback error");
+    
+    // Should clean up localStorage even when callback throws error
+    expect(window.localStorage.getItem("fbui:upgrade:oldUserId")).toBeNull();
+  });
+});
diff --git a/packages/core/src/behaviors/anonymous-upgrade.ts b/packages/core/src/behaviors/anonymous-upgrade.ts
@@ -1,32 +1,66 @@
-import { AuthCredential, AuthProvider, linkWithCredential, linkWithRedirect } from "firebase/auth";
+import { AuthCredential, AuthProvider, linkWithCredential, linkWithRedirect, UserCredential } from "firebase/auth";
 import { FirebaseUIConfiguration } from "~/config";
-import { RedirectHandler } from "./utils";
 import { getBehavior } from "~/behaviors";
 
-export const autoUpgradeAnonymousCredentialHandler = async (ui: FirebaseUIConfiguration, credential: AuthCredential) => {
+export type OnUpgradeCallback = (ui: FirebaseUIConfiguration, oldUserId: string, credential: UserCredential) => Promise<void> | void;
+
+export const autoUpgradeAnonymousCredentialHandler = async (ui: FirebaseUIConfiguration, credential: AuthCredential, onUpgrade?: OnUpgradeCallback) => {
   const currentUser = ui.auth.currentUser;
 
   if (!currentUser?.isAnonymous) {
     return;
   }
 
+  const oldUserId = currentUser.uid;
+
   ui.setState("pending");
   const result = await linkWithCredential(currentUser, credential);
 
+  if (onUpgrade) {
+    await onUpgrade(ui, oldUserId, result);
+  }
+
   ui.setState("idle");
   return result;
 };
 
-export const autoUpgradeAnonymousProviderHandler = async (ui: FirebaseUIConfiguration, provider: AuthProvider) => {
+export const autoUpgradeAnonymousProviderHandler = async (ui: FirebaseUIConfiguration, provider: AuthProvider, onUpgrade?: OnUpgradeCallback) => {
   const currentUser = ui.auth.currentUser;
 
   if (!currentUser?.isAnonymous) {
     return;
   }
 
-  return getBehavior(ui, "providerLinkStrategy")(ui, currentUser, provider);
+  const oldUserId = currentUser.uid;
+
+  window.localStorage.setItem("fbui:upgrade:oldUserId", oldUserId);
+
+  const result = await getBehavior(ui, "providerLinkStrategy")(ui, currentUser, provider);
+
+  // If we got here, the user has been linked via a popup, so we need to call the onUpgrade callback
+  // and delete the oldUserId from localStorage.
+  // If we didn't get here, they'll be redirected and we'll handle the result inside of the autoUpgradeAnonymousUserRedirectHandler.
+
+  window.localStorage.removeItem("fbui:upgrade:oldUserId");
+
+  if (onUpgrade) {
+    await onUpgrade(ui, oldUserId, result);
+  }
+
+  return result;
 };
 
-export const autoUpgradeAnonymousUserRedirectHandler: RedirectHandler = async () => {
-  // TODO
+export const autoUpgradeAnonymousUserRedirectHandler = async (ui: FirebaseUIConfiguration, credential: UserCredential | null, onUpgrade?: OnUpgradeCallback) => {
+  const oldUserId = window.localStorage.getItem("fbui:upgrade:oldUserId");
+
+  // Always clean up localStorage once we've retrieved the oldUserId
+  if (oldUserId) {
+    window.localStorage.removeItem("fbui:upgrade:oldUserId");
+  }
+
+  if (!onUpgrade || !oldUserId || !credential) {
+    return;
+  }
+
+  await onUpgrade(ui, oldUserId, credential);
 };
diff --git a/packages/core/src/behaviors/index.test.ts b/packages/core/src/behaviors/index.test.ts
@@ -9,6 +9,13 @@ import {
   defaultBehaviors,
 } from "./index";
 
+// Mock the anonymous-upgrade handlers
+vi.mock("./anonymous-upgrade", () => ({
+  autoUpgradeAnonymousCredentialHandler: vi.fn(),
+  autoUpgradeAnonymousProviderHandler: vi.fn(),
+  autoUpgradeAnonymousUserRedirectHandler: vi.fn(),
+}));
+
 vi.mock("firebase/auth", () => ({
   RecaptchaVerifier: vi.fn(),
 }));
@@ -125,6 +132,43 @@ describe("autoUpgradeAnonymousUsers", () => {
     expect(typeof behavior.autoUpgradeAnonymousUserRedirectHandler.handler).toBe("function");
   });
 
+  it("should work with onUpgrade callback option", () => {
+    const mockOnUpgrade = vi.fn();
+    const behavior = autoUpgradeAnonymousUsers({ onUpgrade: mockOnUpgrade });
+
+    expect(behavior).toHaveProperty("autoUpgradeAnonymousCredential");
+    expect(behavior).toHaveProperty("autoUpgradeAnonymousProvider");
+    expect(behavior).toHaveProperty("autoUpgradeAnonymousUserRedirectHandler");
+
+    expect(typeof behavior.autoUpgradeAnonymousCredential.handler).toBe("function");
+    expect(typeof behavior.autoUpgradeAnonymousProvider.handler).toBe("function");
+    expect(typeof behavior.autoUpgradeAnonymousUserRedirectHandler.handler).toBe("function");
+  });
+
+  it("should pass onUpgrade callback to handlers when called", async () => {
+    const mockOnUpgrade = vi.fn();
+    const behavior = autoUpgradeAnonymousUsers({ onUpgrade: mockOnUpgrade });
+    
+    const mockUI = createMockUI();
+    const mockCredential = { providerId: "password" } as any;
+    const mockProvider = { providerId: "google.com" } as any;
+    const mockUserCredential = { user: { uid: "upgraded-123" } } as any;
+
+    const { 
+      autoUpgradeAnonymousCredentialHandler,
+      autoUpgradeAnonymousProviderHandler,
+      autoUpgradeAnonymousUserRedirectHandler 
+    } = await import("./anonymous-upgrade");
+
+    await behavior.autoUpgradeAnonymousCredential.handler(mockUI, mockCredential);
+    await behavior.autoUpgradeAnonymousProvider.handler(mockUI, mockProvider);
+    await behavior.autoUpgradeAnonymousUserRedirectHandler.handler(mockUI, mockUserCredential);
+
+    expect(autoUpgradeAnonymousCredentialHandler).toHaveBeenCalledWith(mockUI, mockCredential, mockOnUpgrade);
+    expect(autoUpgradeAnonymousProviderHandler).toHaveBeenCalledWith(mockUI, mockProvider, mockOnUpgrade);
+    expect(autoUpgradeAnonymousUserRedirectHandler).toHaveBeenCalledWith(mockUI, mockUserCredential, mockOnUpgrade);
+  });
+
   it("should not include other behaviors", () => {
     const behavior = autoUpgradeAnonymousUsers();
 
diff --git a/packages/core/src/behaviors/index.ts b/packages/core/src/behaviors/index.ts
@@ -1,5 +1,5 @@
 import type { FirebaseUIConfiguration } from "~/config";
-import type { RecaptchaVerifier } from "firebase/auth";
+import type { RecaptchaVerifier, UserCredential } from "firebase/auth";
 import * as anonymousUpgradeHandlers from "./anonymous-upgrade";
 import * as autoAnonymousLoginHandlers from "./auto-anonymous-login";
 import * as recaptchaHandlers from "./recaptcha";
@@ -21,12 +21,18 @@ type Registry = {
   >;
   autoUpgradeAnonymousProvider: CallableBehavior<typeof anonymousUpgradeHandlers.autoUpgradeAnonymousProviderHandler>;
   autoUpgradeAnonymousUserRedirectHandler: RedirectBehavior<
-    typeof anonymousUpgradeHandlers.autoUpgradeAnonymousUserRedirectHandler
+    (
+      ui: FirebaseUIConfiguration,
+      credential: UserCredential | null,
+      onUpgrade?: anonymousUpgradeHandlers.OnUpgradeCallback
+    ) => ReturnType<typeof anonymousUpgradeHandlers.autoUpgradeAnonymousUserRedirectHandler>
   >;
   recaptchaVerification: CallableBehavior<(ui: FirebaseUIConfiguration, element: HTMLElement) => RecaptchaVerifier>;
   providerSignInStrategy: CallableBehavior<providerStrategyHandlers.ProviderSignInStrategyHandler>;
   providerLinkStrategy: CallableBehavior<providerStrategyHandlers.ProviderLinkStrategyHandler>;
-  oneTapSignIn: InitBehavior<(ui: FirebaseUIConfiguration) => ReturnType<typeof oneTapSignInHandlers.oneTapSignInHandler>>;
+  oneTapSignIn: InitBehavior<
+    (ui: FirebaseUIConfiguration) => ReturnType<typeof oneTapSignInHandlers.oneTapSignInHandler>
+  >;
 };
 
 export type Behavior<T extends keyof Registry = keyof Registry> = Pick<Registry, T>;
@@ -38,14 +44,25 @@ export function autoAnonymousLogin(): Behavior<"autoAnonymousLogin"> {
   };
 }
 
-export function autoUpgradeAnonymousUsers(): Behavior<
+export type AutoUpgradeAnonymousUsersOptions = {
+  onUpgrade?: anonymousUpgradeHandlers.OnUpgradeCallback;
+};
+
+export function autoUpgradeAnonymousUsers(
+  options?: AutoUpgradeAnonymousUsersOptions
+): Behavior<
   "autoUpgradeAnonymousCredential" | "autoUpgradeAnonymousProvider" | "autoUpgradeAnonymousUserRedirectHandler"
 > {
   return {
-    autoUpgradeAnonymousCredential: callableBehavior(anonymousUpgradeHandlers.autoUpgradeAnonymousCredentialHandler),
-    autoUpgradeAnonymousProvider: callableBehavior(anonymousUpgradeHandlers.autoUpgradeAnonymousProviderHandler),
+    autoUpgradeAnonymousCredential: callableBehavior((ui, credential) =>
+      anonymousUpgradeHandlers.autoUpgradeAnonymousCredentialHandler(ui, credential, options?.onUpgrade)
+    ),
+    autoUpgradeAnonymousProvider: callableBehavior((ui, provider) =>
+      anonymousUpgradeHandlers.autoUpgradeAnonymousProviderHandler(ui, provider, options?.onUpgrade)
+    ),
     autoUpgradeAnonymousUserRedirectHandler: redirectBehavior(
-      anonymousUpgradeHandlers.autoUpgradeAnonymousUserRedirectHandler
+      (ui, credential) =>
+        anonymousUpgradeHandlers.autoUpgradeAnonymousUserRedirectHandler(ui, credential, options?.onUpgrade)
     ),
   };
 }
