Merge remote-tracking branch 'origin/main' into 9193-fix-flaky-playwr…

…ight-tests

src/auth/useRequiredPartnerAuth.test.tsx

@@ -54,10 +54,7 @@ beforeEach(() => {
   jest.clearAllMocks();
   // eslint-disable-next-line no-restricted-syntax -- we really do want to resolve to undefined
   getExtensionTokenMock.mockResolvedValue(undefined);
-  useManagedStorageStateMock.mockReturnValue({
-    data: {},
-    isLoading: false,
-  });
+  useManagedStorageStateMock.mockReturnValue(valueToAsyncState({}));
 
   usePartnerAuthDataMock.mockReturnValue(valueToAsyncState(undefined));
 });
@@ -163,10 +160,9 @@ describe("useRequiredPartnerAuth", () => {
   });
 
   test("requires integration for managed storage partner", async () => {
-    useManagedStorageStateMock.mockReturnValue({
-      data: { partnerId: "automation-anywhere" },
-      isLoading: false,
-    });
+    useManagedStorageStateMock.mockReturnValue(
+      valueToAsyncState({ partnerId: "automation-anywhere" }),
+    );
 
     mockAnonymousMeApiResponse();
 

src/extensionConsole/pages/BrowserBanner.test.tsx

@@ -21,9 +21,11 @@ import BrowserBanner from "@/extensionConsole/pages/BrowserBanner";
 import { waitForEffect } from "@/testUtils/testHelpers";
 import { screen } from "@testing-library/react";
 import { INTERNAL_reset } from "@/store/enterprise/managedStorage";
+import { INTERNAL_reset as resetAsyncExternalStore } from "@/hooks/useAsyncExternalStore";
 
 beforeEach(async () => {
   await INTERNAL_reset();
+  resetAsyncExternalStore();
   await browser.storage.managed.clear();
 });
 

src/hooks/useTheme.test.ts

@@ -17,23 +17,33 @@
 
 import useTheme from "@/hooks/useTheme";
 import { renderHook } from "@/pageEditor/testHelpers";
-import useAsyncExternalStore from "@/hooks/useAsyncExternalStore";
 import { initialTheme } from "@/themes/themeStore";
-import { type AsyncState } from "@/types/sliceTypes";
-import { themeStorage } from "@/themes/themeUtils";
+import { type ThemeAssets, themeStorage } from "@/themes/themeUtils";
 import { activateTheme } from "@/background/messenger/api";
-import { readManagedStorageByKey } from "@/store/enterprise/managedStorage";
+import { readManagedStorage } from "@/store/enterprise/managedStorage";
+import { INTERNAL_reset as resetAsyncExternalStore } from "@/hooks/useAsyncExternalStore";
+
+jest.mock("@/themes/themeUtils", () => ({
+  ...jest.requireActual("@/themes/themeUtils"),
+  themeStorage: {
+    get: jest.fn(),
+    onChanged: jest.fn(),
+  },
+}));
+
+jest.mock("@/store/enterprise/managedStorage", () => ({
+  ...jest.requireActual("@/store/enterprise/managedStorage"),
+  readManagedStorage: jest.fn(),
+}));
 
 afterEach(() => {
   jest.clearAllMocks();
 });
 
-jest.mock("@/hooks/useAsyncExternalStore");
 jest.mock("@/background/messenger/api");
-jest.mock("@/store/enterprise/managedStorage");
 
-const customTheme = {
-  themeName: "custom",
+const customTheme: ThemeAssets = {
+  themeName: "default",
   showSidebarLogo: true,
   customSidebarLogo: "https://example.com/custom-logo.png",
   toolbarIcon: "https://example.com/custom-icon.svg",
@@ -46,31 +56,33 @@ const customTheme = {
 
 describe("useTheme", () => {
   beforeEach(async () => {
-    jest
-      .mocked(useAsyncExternalStore)
-      .mockReturnValue({ data: initialTheme, isLoading: false } as AsyncState);
-    // eslint-disable-next-line no-restricted-syntax -- this func requires a parameter
-    jest.mocked(readManagedStorageByKey).mockResolvedValue(undefined);
+    resetAsyncExternalStore();
+
+    jest.mocked(themeStorage.get).mockResolvedValue({
+      ...initialTheme,
+      lastFetched: Date.now(),
+    });
+    jest.mocked(readManagedStorage).mockResolvedValue({});
+  });
+
+  afterEach(() => {
+    jest.useRealTimers();
   });
 
-  test("calls useAsyncExternalStore and gets current theme state", async () => {
+  test("calls themeStorage to get the current theme state", async () => {
     const { result: themeResult, waitForNextUpdate } = renderHook(() =>
       useTheme(),
     );
 
     await waitForNextUpdate();
 
-    expect(useAsyncExternalStore).toHaveBeenNthCalledWith(
-      1,
-      expect.any(Function),
-      themeStorage.get,
-    );
+    expect(themeStorage.get).toHaveBeenCalledOnce();
 
-    expect(themeResult.current).toStrictEqual({
+    expect(themeResult.current).toMatchObject({
       activeTheme: {
         themeName: "default",
         customSidebarLogo: null,
-        lastFetched: null,
+        lastFetched: expect.any(Number),
         logo: { regular: "test-file-stub", small: "test-file-stub" },
         showSidebarLogo: true,
         toolbarIcon: null,
@@ -81,20 +93,17 @@ describe("useTheme", () => {
 
   it("calls activateTheme after loading is done and it hasn't been called recently", async () => {
     jest.useFakeTimers();
+    const { rerender, waitForNextUpdate } = renderHook(() => useTheme());
+    await waitForNextUpdate();
 
-    jest.mocked(useAsyncExternalStore).mockReturnValue({
-      data: { ...initialTheme, lastFetched: Date.now() },
-      isLoading: false,
-    } as AsyncState);
-
-    let result = renderHook(() => useTheme());
-    await result.waitForNextUpdate();
     expect(activateTheme).not.toHaveBeenCalled();
 
     jest.advanceTimersByTime(125_000);
+    resetAsyncExternalStore();
+
+    rerender();
+    await waitForNextUpdate();
 
-    result = renderHook(() => useTheme());
-    await result.waitForNextUpdate();
     expect(activateTheme).toHaveBeenCalledOnce();
   });
 
@@ -108,36 +117,51 @@ describe("useTheme", () => {
   ])(
     "handles showSidebarLogo policy (policy: $policyValue, theme: $themeValue, expected: $expectedValue)",
     async ({ policyValue, themeValue, expectedValue }) => {
-      jest.mocked(useAsyncExternalStore).mockReturnValue({
-        data: { ...customTheme, showSidebarLogo: themeValue },
-        isLoading: false,
-      } as AsyncState);
-      jest.mocked(readManagedStorageByKey).mockResolvedValue(policyValue);
+      const lastFetched = Date.now();
+      jest.mocked(themeStorage.get).mockResolvedValue({
+        ...customTheme,
+        showSidebarLogo: themeValue,
+        lastFetched,
+      });
+
+      jest.mocked(readManagedStorage).mockResolvedValue({
+        showSidebarLogo: policyValue,
+      });
 
       const { result, waitForNextUpdate } = renderHook(() => useTheme());
 
       await waitForNextUpdate();
 
-      expect(result.current.activeTheme.showSidebarLogo).toBe(expectedValue);
+      expect(result.current.activeTheme).toMatchObject({
+        ...customTheme,
+        lastFetched,
+        showSidebarLogo: expectedValue,
+      });
     },
   );
 
-  it("uses activeTheme when an error occurs in managed storage", async () => {
-    jest.mocked(useAsyncExternalStore).mockReturnValue({
-      data: customTheme,
-      isLoading: false,
-    } as AsyncState);
+  it.each([{ showSidebarLogo: true }, { showSidebarLogo: false }])(
+    "uses activeTheme when an error occurs in managed storage (showSidebarLogo: $showSidebarLogo)",
+    async ({ showSidebarLogo }) => {
+      const customThemeWithSidebarLogo = {
+        ...customTheme,
+        showSidebarLogo,
+        lastFetched: Date.now(),
+      };
 
-    jest
-      .mocked(readManagedStorageByKey)
-      .mockRejectedValue(new Error("Managed storage error"));
+      jest
+        .mocked(themeStorage.get)
+        .mockResolvedValue(customThemeWithSidebarLogo);
 
-    const { result, waitForNextUpdate } = renderHook(() => useTheme());
+      jest
+        .mocked(readManagedStorage)
+        .mockRejectedValue(new Error("Managed storage error"));
 
-    await waitForNextUpdate();
+      const { result, waitForNextUpdate } = renderHook(() => useTheme());
 
-    expect(result.current.activeTheme.showSidebarLogo).toBe(
-      customTheme.showSidebarLogo,
-    );
-  });
+      await waitForNextUpdate();
+
+      expect(result.current.activeTheme.showSidebarLogo).toBe(showSidebarLogo);
+    },
+  );
 });

src/hooks/useTheme.ts

@@ -25,8 +25,7 @@ import {
 import { initialTheme } from "@/themes/themeStore";
 import useAsyncExternalStore from "@/hooks/useAsyncExternalStore";
 import { activateTheme } from "@/background/messenger/api";
-import useAsyncState from "@/hooks/useAsyncState";
-import { readManagedStorageByKey } from "@/store/enterprise/managedStorage";
+import useManagedStorageState from "@/store/enterprise/useManagedStorageState";
 
 const themeStorageSubscribe = (callback: () => void) => {
   const abortController = new AbortController();
@@ -47,8 +46,10 @@ function useTheme(): { activeTheme: ThemeAssets; isLoading: boolean } {
   const { data: cachedTheme, isLoading: isCachedThemeLoading } =
     useAsyncExternalStore(themeStorageSubscribe, themeStorage.get);
 
-  const { data: showSidebarLogoOverride, isLoading: isManagedStorageLoading } =
-    useAsyncState(async () => readManagedStorageByKey("showSidebarLogo"), []);
+  const { data: managedStorageState, isLoading: isManagedStorageLoading } =
+    useManagedStorageState();
+
+  const showSidebarLogoOverride = managedStorageState?.showSidebarLogo;
 
   const isLoading = isManagedStorageLoading || isCachedThemeLoading;
 

src/sandbox/postMessage.test.ts

@@ -24,7 +24,6 @@ import postMessage, {
   addPostMessageListener,
   pendingMessageMetadataMap,
   type RequestPacket,
-  SandboxTimeoutError,
 } from "./postMessage";
 import { sleep } from "@/utils/timeUtils";
 
@@ -175,18 +174,14 @@ describe("SandboxTimeoutError", () => {
 
     jest.runAllTimers();
 
-    await expect(promise).rejects.toThrow(SandboxTimeoutError);
-    await expect(promise).rejects.toMatchObject({
-      name: "SandboxTimeoutError",
-      message:
-        "Message SANDBOX_PING did not receive a response within 5 seconds",
-      sandboxMessage: {
-        type: "SANDBOX_PING",
-        payloadSize: expect.any(Number),
-        timestamp: expect.any(Number),
-      },
-      pendingSandboxMessages: [],
-    });
+    await expect(promise).rejects.toThrow(
+      expect.objectContaining({
+        name: "SandboxTimeoutError",
+        message: expect.stringMatching(
+          /Sandbox message timed out: type=SANDBOX_PING, sent=\d+, payloadSize=\d+B, pendingMessageCount=0/,
+        ),
+      }),
+    );
   });
 
   test("throws SandboxTimeoutError with correct information for multiple messages", async () => {
@@ -216,62 +211,34 @@ describe("SandboxTimeoutError", () => {
       recipient: unresponsiveChannel as Window,
     });
 
-    await expect(promise1).rejects.toThrow(SandboxTimeoutError);
-    await expect(promise1).rejects.toMatchObject({
-      name: "SandboxTimeoutError",
-      message:
-        "Message SANDBOX_FOO did not receive a response within 5 seconds",
-      sandboxMessage: {
-        type: "SANDBOX_FOO",
-        payloadSize: expect.any(Number),
-        timestamp: expect.any(Number),
-      },
-      pendingSandboxMessages: expect.arrayContaining([
-        {
-          type: "SANDBOX_BAR",
-          payloadSize: expect.any(Number),
-          timestamp: expect.any(Number),
-        },
-        {
-          type: "SANDBOX_BAZ",
-          payloadSize: expect.any(Number),
-          timestamp: expect.any(Number),
-        },
-      ]),
-    });
+    await expect(promise1).rejects.toThrow(
+      expect.objectContaining({
+        name: "SandboxTimeoutError",
+        message: expect.stringMatching(
+          /Sandbox message timed out: type=SANDBOX_FOO, sent=\d+, payloadSize=\d+B, pendingMessageCount=2/,
+        ),
+      }),
+    );
 
-    await expect(promise2).rejects.toThrow(SandboxTimeoutError);
-    await expect(promise2).rejects.toMatchObject({
-      name: "SandboxTimeoutError",
-      message:
-        "Message SANDBOX_BAR did not receive a response within 5 seconds",
-      sandboxMessage: {
-        type: "SANDBOX_BAR",
-        payloadSize: expect.any(Number),
-        timestamp: expect.any(Number),
-      },
-      pendingSandboxMessages: expect.arrayContaining([
-        {
-          type: "SANDBOX_BAZ",
-          payloadSize: expect.any(Number),
-          timestamp: expect.any(Number),
-        },
-      ]),
-    });
+    await expect(promise2).rejects.toThrow(
+      expect.objectContaining({
+        name: "SandboxTimeoutError",
+        message: expect.stringMatching(
+          /Sandbox message timed out: type=SANDBOX_BAR, sent=\d+, payloadSize=\d+B, pendingMessageCount=1/,
+        ),
+      }),
+    );
 
     jest.runAllTimers();
-    await expect(promise3).rejects.toThrow(SandboxTimeoutError);
-    await expect(promise3).rejects.toMatchObject({
-      name: "SandboxTimeoutError",
-      message:
-        "Message SANDBOX_BAZ did not receive a response within 5 seconds",
-      sandboxMessage: {
-        type: "SANDBOX_BAZ",
-        payloadSize: expect.any(Number),
-        timestamp: expect.any(Number),
-      },
-      pendingSandboxMessages: [],
-    });
+
+    await expect(promise3).rejects.toThrow(
+      expect.objectContaining({
+        name: "SandboxTimeoutError",
+        message: expect.stringMatching(
+          /Sandbox message timed out: type=SANDBOX_BAZ, sent=\d+, payloadSize=\d+B, pendingMessageCount=0/,
+        ),
+      }),
+    );
   });
 
   test("clears metadata for resolved messages", async () => {