#9227 Refactor useTheme to use useManagedStorageState

src/auth/useRequiredPartnerAuth.test.tsx

@@ -39,6 +39,8 @@ import usePartnerAuthData from "@/auth/usePartnerAuthData";
 import { Milestones } from "@/data/model/UserMilestone";
 import { getDeploymentKey } from "@/auth/deploymentKey";
 import { getExtensionToken } from "@/auth/authStorage";
+import type { AsyncState } from "@/types/sliceTypes";
+import { type ManagedStorageState } from "@/store/enterprise/managedStorageTypes";
 
 jest.mock("@/store/enterprise/useManagedStorageState");
 jest.mock("@/auth/usePartnerAuthData");
@@ -57,7 +59,7 @@ beforeEach(() => {
   useManagedStorageStateMock.mockReturnValue({
     data: {},
     isLoading: false,
-  });
+  } as AsyncState<ManagedStorageState>);
 
   usePartnerAuthDataMock.mockReturnValue(valueToAsyncState(undefined));
 });
@@ -166,7 +168,7 @@ describe("useRequiredPartnerAuth", () => {
     useManagedStorageStateMock.mockReturnValue({
       data: { partnerId: "automation-anywhere" },
       isLoading: false,
-    });
+    } as AsyncState<ManagedStorageState>);
 
     mockAnonymousMeApiResponse();
 

src/hooks/useTheme.test.ts

@@ -17,23 +17,32 @@
 
 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";
+
+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 +55,31 @@
 
 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);
+    jest.mocked(themeStorage.get).mockResolvedValue({
+      ...initialTheme,
+      lastFetched: Date.now(),
+    });
+    jest.mocked(readManagedStorage).mockResolvedValue({});
   });
 
-  test("calls useAsyncExternalStore and gets current theme state", async () => {
+  afterEach(() => {
+    jest.useRealTimers();
+  });
+
+  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 +90,14 @@
 
   it("calls activateTheme after loading is done and it hasn't been called recently", async () => {
     jest.useFakeTimers();
+    renderHook(() => useTheme());
 
-    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);
 
-    result = renderHook(() => useTheme());
-    await result.waitForNextUpdate();
+    renderHook(() => useTheme());
+
     expect(activateTheme).toHaveBeenCalledOnce();
   });
 
@@ -108,36 +111,48 @@
   ])(
     "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);
+      jest.mocked(themeStorage.get).mockResolvedValue({
+        ...customTheme,
+        showSidebarLogo: themeValue,
+        lastFetched: Date.now(),
+      });
+
+      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: expect.any(Number),
+        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 } = renderHook(() => useTheme());
 
-    expect(result.current.activeTheme.showSidebarLogo).toBe(
-      customTheme.showSidebarLogo,
-    );
-  });
+      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/useManagedStorageState.test.ts

@@ -25,11 +25,17 @@
 });
 
 describe("useManagedStorageState", () => {
   it("waits on uninitialized state", async () => {
     const { result, waitFor } = renderHook(() => useManagedStorageState());
     expect(result.current).toStrictEqual({
+      currentData: undefined,
       data: undefined,
+      error: undefined,
+      isError: false,
+      isFetching: true,
       isLoading: true,
+      isSuccess: false,
+      isUninitialized: false,
     });
 
     await waitFor(
@@ -52,13 +58,22 @@
       useManagedStorageState(),
     );
 
     await waitForNextUpdate();
 
     expect(result.current).toStrictEqual({
+      currentData: {
+        partnerId: "taco-bell",
+      },
       data: {
         partnerId: "taco-bell",
       },
+      error: undefined,
+      isError: false,
+      isFetching: false,
       isLoading: false,
+      isSuccess: true,
+      isUninitialized: false,
+      refetch: expect.any(Function),
     });
   });
 

src/store/enterprise/useManagedStorageState.ts

@@ -15,21 +15,12 @@
  * along with this program.  If not, see <http://www.gnu.org/licenses/>.
  */
 
-import { useSyncExternalStore } from "use-sync-external-store/shim";
 import {
-  getSnapshot,
-  initManagedStorage,
   managedStorageStateChange,
+  readManagedStorage,
 } from "@/store/enterprise/managedStorage";
-import { useEffect } from "react";
-import type { ManagedStorageState } from "@/store/enterprise/managedStorageTypes";
-import type { Nullishable } from "@/utils/nullishUtils";
 import { expectContext } from "@/utils/expectContext";
-
-type HookState = {
-  data: Nullishable<ManagedStorageState>;
-  isLoading: boolean;
-};
+import useAsyncExternalStore from "@/hooks/useAsyncExternalStore";
 
 // NOTE: can't share subscribe methods across generators currently for useAsyncExternalStore because it maintains
 // a map of subscriptions to state controllers. See https://github.com/pixiebrix/pixiebrix-extension/issues/7789
@@ -46,18 +37,8 @@ function subscribe(callback: () => void): () => void {
 /**
  * React hook to get the current state of managed storage.
  */
-function useManagedStorageState(): HookState {
-  useEffect(() => {
-    // `initManagedStorage` is wrapped in once, so safe to call from multiple locations in the tree.
-    void initManagedStorage();
-  }, []);
-
-  const data = useSyncExternalStore(subscribe, getSnapshot);
-
-  return {
-    data,
-    isLoading: data == null,
-  };
+function useManagedStorageState() {
+  return useAsyncExternalStore(subscribe, readManagedStorage);
 }
 
 export default useManagedStorageState;