#9227 Refactor useTheme to use useManagedStorageState

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/store/enterprise/managedStorage.ts

@@ -318,17 +318,6 @@ export function isInitialized(): boolean {
   return managedStorageSnapshot != null;
 }
 
-/**
- * Get a _synchronous_ snapshot of the managed storage state.
- * @see useManagedStorageState
- * @see readManagedStorage
- */
-export function getSnapshot(): Nullishable<ManagedStorageState> {
-  expectContext("extension");
-
-  return managedStorageSnapshot;
-}
-
 /**
  * Helper method for resetting the module for testing.
  */

src/store/enterprise/useManagedStorageState.test.ts

@@ -15,64 +15,59 @@
  * along with this program.  If not, see <http://www.gnu.org/licenses/>.
  */
 
-import { INTERNAL_reset } from "@/store/enterprise/managedStorage";
+import { INTERNAL_reset as resetManagedStorage } from "@/store/enterprise/managedStorage";
+import { INTERNAL_reset as resetAsyncExternalStore } from "@/hooks/useAsyncExternalStore";
 import { renderHook } from "@testing-library/react-hooks";
 import useManagedStorageState from "@/store/enterprise/useManagedStorageState";
+import {
+  loadingAsyncStateFactory,
+  valueToAsyncState,
+} from "@/utils/asyncStateUtils";
 
 beforeEach(async () => {
-  await INTERNAL_reset();
+  await resetManagedStorage();
+  resetAsyncExternalStore();
   await browser.storage.managed.clear();
 });
 
 describe("useManagedStorageState", () => {
   it("waits on uninitialized state", async () => {
-    const { result, waitFor } = renderHook(() => useManagedStorageState());
-    expect(result.current).toStrictEqual({
-      data: undefined,
-      isLoading: true,
+    const { result, waitForNextUpdate } = renderHook(() =>
+      useManagedStorageState(),
+    );
+    expect(result.current).toStrictEqual(loadingAsyncStateFactory());
+
+    await waitForNextUpdate({
+      // `readManagedStorage` will take a few seconds if there are no policies set
+      timeout: 5000,
     });
 
-    await waitFor(
-      () => {
-        expect(result.current).toStrictEqual({
-          data: {},
-          isLoading: false,
-        });
-      },
-      // XXX: figure out how to use fake timers to avoid slowing down test suite
-      // Must be longer than MAX_MANAGED_STORAGE_WAIT_MILLIS
-      { timeout: 5000 },
-    );
+    expect(result.current).toStrictEqual(valueToAsyncState({}));
   });
 
   it("handles already initialized state", async () => {
-    await browser.storage.managed.set({ partnerId: "taco-bell" });
+    const expectedPolicy = { partnerId: "foo" };
+    await browser.storage.managed.set(expectedPolicy);
 
     const { result, waitForNextUpdate } = renderHook(() =>
       useManagedStorageState(),
     );
 
     await waitForNextUpdate();
 
-    expect(result.current).toStrictEqual({
-      data: {
-        partnerId: "taco-bell",
-      },
-      isLoading: false,
-    });
+    expect(result.current).toStrictEqual(valueToAsyncState(expectedPolicy));
   });
 
-  // Can't test because mock doesn't fire change events: https://github.com/clarkbw/jest-webextension-mock/issues/170
-  it.skip("listens for changes", async () => {
+  it("listens for changes", async () => {
     const { result, waitForNextUpdate } = renderHook(() =>
       useManagedStorageState(),
     );
     expect(result.current.data).toBeUndefined();
-    await browser.storage.managed.set({ partnerId: "taco-bell" });
+
+    const expectedPolicy = { partnerId: "foo" };
+    await browser.storage.managed.set(expectedPolicy);
+
     await waitForNextUpdate();
-    expect(result.current).toStrictEqual({
-      data: { partnerId: "taco-bell" },
-      isLoading: false,
-    });
+    expect(result.current).toStrictEqual(valueToAsyncState(expectedPolicy));
   });
 });