DiamondLightSource · VictoriaBeilsten-Edmands · Jul 3, 2026 · Jul 3, 2026
diff --git a/dev/index.html b/dev/index.html
@@ -6,6 +6,14 @@
     <title>SciReactUI Dev</title>
   </head>
   <body>
+    <script>
+      (function () {
+        var mode = localStorage.getItem("mui-mode");
+        if (mode) {
+          document.documentElement.setAttribute("data-mode", mode);
+        }
+      })();
+    </script>
     <div id="root"></div>
     <script type="module" src="/src/main.tsx"></script>
   </body>

diff --git a/src/components/controls/ColourSchemeButton.test.tsx b/src/components/controls/ColourSchemeButton.test.tsx
@@ -1,37 +1,50 @@
 import { fireEvent } from "@testing-library/react";
+import { useColorScheme } from "@mui/material/styles";
 
 import { ColourSchemeButton } from "./ColourSchemeButton";
-import { ColourSchemes } from "../../utils/globals";
 import { renderWithProviders } from "../../__test-utils__/helpers";
 
 const mockSetColorScheme = vi.fn();
-vi.mock("@mui/material", async () => {
-  const actualEnums = await vi.importActual("../../utils/globals");
+vi.mock("@mui/material/styles", async () => {
+  const actual = await vi.importActual<typeof import("@mui/material/styles")>(
+    "@mui/material/styles",
+  );
 
   return {
-    ...(await vi.importActual("@mui/material")),
-    useColorScheme: vi.fn().mockReturnValue({
-      // @ts-expect-error module doesn't have a type
-      colorScheme: actualEnums.ColourSchemes.Dark,
-      setColorScheme: (scheme: ColourSchemes) => mockSetColorScheme(scheme),
-    }),
+    ...actual,
+    useColorScheme: vi.fn(),
   };
 });
 
+const mockUseColorScheme = vi.mocked(useColorScheme);
+
 describe("ColourSchemeButton", () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+
+    mockUseColorScheme.mockReturnValue({
+      mode: "dark",
+      systemMode: undefined,
+      setMode: mockSetColorScheme,
+      colorScheme: "dark",
+      allColorSchemes: ["light", "dark"],
+      setColorScheme: vi.fn(),
+    });
+  });
+
   it("should render without errors", () => {
     renderWithProviders(<ColourSchemeButton />);
   });
 
-  it("should show dark icon and button", () => {
+  it("should show button and light mode icon when current mode is dark", () => {
     const { getByTestId, getByRole } = renderWithProviders(
       <ColourSchemeButton />,
     );
 
     const button = getByRole("button");
     expect(button).toBeInTheDocument();
 
-    const icon = getByTestId("BedtimeIcon");
+    const icon = getByTestId("LightModeIcon");
     expect(icon).toBeInTheDocument();
   });
 
@@ -41,7 +54,7 @@ describe("ColourSchemeButton", () => {
     const button = getByRole("button");
     fireEvent.click(button);
 
-    //expect(mockSetColorScheme).toHaveBeenCalledWith(ColourSchemes.Light);
+    expect(mockSetColorScheme).toHaveBeenCalledWith("light");
   });
 
   it("should call local onclick when button clicked", () => {
@@ -55,4 +68,19 @@ describe("ColourSchemeButton", () => {
 
     expect(mockOnClick).toHaveBeenCalled();
   });
+
+  it("should not render while colour scheme is unresolved", () => {
+    mockUseColorScheme.mockReturnValue({
+      mode: undefined,
+      systemMode: undefined,
+      setMode: mockSetColorScheme,
+      colorScheme: undefined,
+      allColorSchemes: ["light", "dark"],
+      setColorScheme: vi.fn(),
+    });
+
+    const { queryByRole } = renderWithProviders(<ColourSchemeButton />);
+
+    expect(queryByRole("button")).not.toBeInTheDocument();
+  });
 });
diff --git a/src/components/controls/ColourSchemeButton.tsx b/src/components/controls/ColourSchemeButton.tsx
@@ -9,6 +9,10 @@ export const ColourSchemeButton = (props: IconButtonProps) => {
   const resolvedMode = mode === "system" ? systemMode : mode;
   const isDark = resolvedMode === "dark";
 
+  if (resolvedMode === undefined) {
+    return null;
+  }
+
   return (
     <IconButton
       aria-label={`Colour scheme switcher: ${resolvedMode ?? "unknown"}`}

diff --git a/src/storybook/Installation.mdx b/src/storybook/Installation.mdx
@@ -161,35 +161,67 @@ Use this page to install SciReactUI, configure the Diamond Design System theme,
 
 </div>
 
-  <div className="sb-container">
-    <div className='sb-section-title'>
-    ### Avoiding Dark Mode Flicker in SSR Apps
-
-    To prevent a flash of incorrect theme (e.g. light mode before switching to dark), include the `InitColorSchemeScript` in your app code.
-
-    This ensures the correct colour scheme is applied immediately, based on system preference or saved user choice.
-
-
-      ```tsx filename="main.tsx"
-      import { InitColorSchemeScript } from "@mui/material";
-
-      function Root() {
-        return (
-          <>
-            <InitColorSchemeScript attribute="[data-mode='%s']" />
-            <App />
-          </>
-        );
-      }
-      ```
-
-      If your app renders into `index.html`, place the script at the top of the body so the mode is set before React mounts.
-
-      This helps avoid the brief light-mode flash when the user prefers dark mode.
-
-      For more info see MUI's documentation on [InitColorSchemeScript](https://v7.mui.com/material-ui/react-init-color-scheme-script/).
-
-    </div>
+<div className="sb-container">
+  <div className="sb-section-title">
+
+    ### Avoiding Dark Mode Flicker
+
+    Applications that support dark mode may briefly render the default colour scheme before the user's preferred mode is applied.
+
+    #### SSR Applications
+
+    To prevent a flash of incorrect theme in SSR applications, use MUI's `InitColorSchemeScript` and place it before your application content so the correct colour scheme is applied before hydration.
+    For example, for Next.js App Router:
+    ```tsx filename="layout.tsx"
+    import InitColorSchemeScript from '@mui/material/InitColorSchemeScript';
+
+    export default function RootLayout(props: { children: React.ReactNode }) {
+    return (
+        <html lang="en" suppressHydrationWarning>
+        <body>
+            <InitColorSchemeScript attribute="data" />
+            {props.children}
+        </body>
+        </html>
+    );
+    }
+    ```
+    If your app renders into `index.html`, place the script at the top of the body so the mode is set before React mounts.
+    For more info see MUI's documentation on [InitColorSchemeScript](https://v7.mui.com/material-ui/react-init-color-scheme-script/).
+
+    #### Client-Side Applications (Vite etc.)
+
+    For client-side applications, initialise the mode before React mounts:
+    ```
+    <body>
+        <script>
+        (function () {
+            var mode = localStorage.getItem("mui-mode");
+            if (mode) {
+            document.documentElement.setAttribute("data-mode", mode);
+            }
+        })();
+        </script>
+        <div id="root"></div>
+        <script type="module" src="/src/main.tsx"></script>
+    </body>
+    ```
+
+    This ensures the correct DiamondDS token set is applied before the first paint, avoiding a flash of the default colour scheme when a saved preference exists.
+
+    DiamondDS uses:
+    ```
+    <html data-mode="light">
+    ```
+    or
+    ```
+    <html data-mode="dark">
+    ```
+    and should be configured with
+    ```
+    colorSchemeSelector: '[data-mode="%s"]'
+    ```
 
   </div>
+  </div>
 </div>