🤖 fix: make Storybook theme deterministic via forcedTheme

Author: ammar-agent

.storybook/preview.tsx

@@ -6,12 +6,8 @@ import {
   useTheme,
 } from "../src/browser/contexts/ThemeContext";
 import isChromatic from "chromatic/isChromatic";
-import { UI_THEME_KEY } from "@/common/constants/storage";
-import { updatePersistedState } from "@/browser/hooks/usePersistedState";
 import "../src/browser/styles/globals.css";
 
-const DEFAULT_THEME: ThemeMode = "light";
-
 const isAutomatedChromaticRun = () => {
   if (typeof window === "undefined" || !isChromatic()) {
     return false;
@@ -36,35 +32,6 @@ const isE2ETestEnv = () => {
   return false;
 };
 
-const getPersistedTheme = (): ThemeMode | undefined => {
-  if (typeof window === "undefined") {
-    return undefined;
-  }
-
-  try {
-    const stored = window.localStorage.getItem(UI_THEME_KEY);
-    return stored ? (JSON.parse(stored) as ThemeMode) : undefined;
-  } catch (error) {
-    console.warn("Failed to read Storybook theme:", error);
-    return undefined;
-  }
-};
-
-const syncStorybookTheme = (mode?: ThemeMode): ThemeMode => {
-  if (typeof window === "undefined") {
-    return mode ?? DEFAULT_THEME; // Keep default aligned with baseline to avoid Chromatic diffs
-  }
-
-  const persisted = getPersistedTheme();
-  const resolved = mode ?? persisted ?? DEFAULT_THEME; // Default to light for Latest when unset
-
-  if (resolved !== persisted) {
-    updatePersistedState(UI_THEME_KEY, resolved);
-  }
-
-  return resolved;
-};
-
 const StorybookThemeToggle: React.FC = () => {
   const { theme, toggleTheme } = useTheme();
 
@@ -114,9 +81,8 @@ const preview: Preview = {
   decorators: [
     (Story, context) => {
       const mode = context.globals.theme as ThemeMode | undefined;
-      syncStorybookTheme(mode);
       return (
-        <ThemeProvider>
+        <ThemeProvider forcedTheme={mode}>
           <Story />
           {!mode && !isE2ETestEnv() && <StorybookThemeToggle />}
         </ThemeProvider>

src/browser/contexts/ThemeContext.test.tsx

@@ -0,0 +1,96 @@
+import { GlobalWindow } from "happy-dom";
+
+// Setup basic DOM environment for testing-library
+const dom = new GlobalWindow();
+(global as any).window = dom.window;
+(global as any).document = dom.window.document;
+// Polyfill console since happy-dom might interfere or we just want standard console
+(global as any).console = console;
+
+import { afterEach, describe, expect, mock, test, beforeEach } from "bun:test";
+
+import { render, cleanup } from "@testing-library/react";
+import React from "react";
+import { ThemeProvider, useTheme } from "./ThemeContext";
+import { UI_THEME_KEY } from "@/common/constants/storage";
+
+// Helper to access internals
+const TestComponent = () => {
+  const { theme, toggleTheme } = useTheme();
+  return (
+    <div>
+      <span data-testid="theme-value">{theme}</span>
+      <button onClick={toggleTheme} data-testid="toggle-btn">Toggle</button>
+    </div>
+  );
+};
+
+describe("ThemeContext", () => {
+  // Mock matchMedia
+  const mockMatchMedia = mock(() => ({
+    matches: false,
+    media: "",
+    onchange: null,
+    addListener: () => {},
+    removeListener: () => {},
+    addEventListener: () => {},
+    removeEventListener: () => {},
+    dispatchEvent: () => true,
+  }));
+
+  beforeEach(() => {
+    // Ensure window exists (Bun test with happy-dom should provide it)
+    if (typeof window !== "undefined") {
+      window.matchMedia = mockMatchMedia;
+      window.localStorage.clear();
+    }
+  });
+
+  afterEach(() => {
+    cleanup();
+    if (typeof window !== "undefined") {
+        window.localStorage.clear();
+    }
+  });
+
+  test("uses persisted state by default", () => {
+    const { getByTestId } = render(
+      <ThemeProvider>
+        <TestComponent />
+      </ThemeProvider>
+    );
+    // If matchMedia matches is false (default mock), resolveSystemTheme returns 'dark' (since it checks prefers-color-scheme: light)
+    // resolveSystemTheme logic: window.matchMedia("(prefers-color-scheme: light)").matches ? "light" : "dark"
+    expect(getByTestId("theme-value").textContent).toBe("dark");
+  });
+
+  test("respects forcedTheme prop", () => {
+    const { getByTestId, rerender } = render(
+      <ThemeProvider forcedTheme="light">
+        <TestComponent />
+      </ThemeProvider>
+    );
+    expect(getByTestId("theme-value").textContent).toBe("light");
+
+    rerender(
+      <ThemeProvider forcedTheme="dark">
+        <TestComponent />
+      </ThemeProvider>
+    );
+    expect(getByTestId("theme-value").textContent).toBe("dark");
+  });
+
+  test("forcedTheme overrides persisted state", () => {
+     window.localStorage.setItem(UI_THEME_KEY, JSON.stringify("light"));
+     
+     const { getByTestId } = render(
+      <ThemeProvider forcedTheme="dark">
+        <TestComponent />
+      </ThemeProvider>
+    );
+    expect(getByTestId("theme-value").textContent).toBe("dark");
+    
+    // Check that localStorage is still light (since forcedTheme doesn't write to storage by itself)
+    expect(JSON.parse(window.localStorage.getItem(UI_THEME_KEY)!)).toBe("light");
+  });
+});

src/browser/contexts/ThemeContext.tsx

@@ -51,11 +51,19 @@ function applyThemeToDocument(theme: ThemeMode) {
   }
 }
 
-export function ThemeProvider(props: { children: ReactNode }) {
-  const [theme, setTheme] = usePersistedState<ThemeMode>(UI_THEME_KEY, resolveSystemTheme(), {
+export function ThemeProvider({
+  children,
+  forcedTheme,
+}: {
+  children: ReactNode;
+  forcedTheme?: ThemeMode;
+}) {
+  const [persistedTheme, setTheme] = usePersistedState<ThemeMode>(UI_THEME_KEY, resolveSystemTheme(), {
     listener: true,
   });
 
+  const theme = forcedTheme ?? persistedTheme;
+
   useLayoutEffect(() => {
     applyThemeToDocument(theme);
   }, [theme]);
@@ -73,7 +81,7 @@ export function ThemeProvider(props: { children: ReactNode }) {
     [setTheme, theme, toggleTheme]
   );
 
-  return <ThemeContext.Provider value={value}>{props.children}</ThemeContext.Provider>;
+  return <ThemeContext.Provider value={value}>{children}</ThemeContext.Provider>;
 }
 
 export function useTheme(): ThemeContextValue {