mui · siriwatknp · Aug 24, 2022 · Aug 9, 2022 · Aug 15, 2022 · Aug 15, 2022
diff --git a/packages/mui-system/src/cssVars/createCssVarsProvider.test.js b/packages/mui-system/src/cssVars/createCssVarsProvider.test.js
@@ -611,7 +611,6 @@ describe('createCssVarsProvider', () => {
       );
 
       expect(screen.getByTestId('current-mode').textContent).to.equal('dark');
-      expect(global.localStorage.setItem.calledWith(customModeStorageKey, 'dark')).to.equal(true);
     });
 
     it('support custom storage window', () => {

diff --git a/packages/mui-system/src/cssVars/useCurrentColorScheme.test.js b/packages/mui-system/src/cssVars/useCurrentColorScheme.test.js
@@ -324,6 +324,73 @@ describe('useCurrentColorScheme', () => {
       });
     });
 
+    it('change only the mode specified as key', () => {
+      const Data = () => {
+        const { setColorScheme, ...data } = useCurrentColorScheme({
+          defaultMode: 'light',
+          defaultLightColorScheme: 'light',
+          defaultDarkColorScheme: 'dark',
+          supportedColorSchemes: ['light', 'paper', 'dark', 'dim'],
+        });
+        return (
+          <div>
+            <div data-testid="data">{JSON.stringify(data)}</div>
+            <button onClick={() => setColorScheme('paper')}>first</button>
+            <button onClick={() => setColorScheme({ dark: 'dim' })}>second</button>
+          </div>
+        );
+      };
+      const { getByText, getByTestId } = render(<Data />);
+
+      fireEvent.click(getByText('first'));
+
+      expect(JSON.parse(getByTestId('data').textContent)).to.deep.equal({
+        mode: 'light',
+        lightColorScheme: 'paper',
+        darkColorScheme: 'dark',
+        colorScheme: 'paper',
+      });
+
+      fireEvent.click(getByText('second'));
+
+      expect(JSON.parse(getByTestId('data').textContent)).to.deep.equal({
+        mode: 'light',
+        lightColorScheme: 'paper',
+        darkColorScheme: 'dim',
+        colorScheme: 'paper',
+      });
+    });
+
+    it('able to setMode and setColorScheme in the same event', () => {
+      const Data = () => {
+        const { setColorScheme, setMode, ...data } = useCurrentColorScheme({
+          defaultLightColorScheme: 'light',
+          defaultDarkColorScheme: 'dark',
+          supportedColorSchemes: ['light', 'paper', 'dark', 'dim'],
+        });
+        return (
+          <button
+            onClick={() => {
+              setMode('dark');
+              setColorScheme({ light: 'paper', dark: 'dim' });
+            }}
+          >
+            {JSON.stringify(data)}
+          </button>
+        );
+      };
+      const { container } = render(<Data />);
+
+      fireEvent.click(container.firstChild);
+
+      expect(JSON.parse(container.firstChild.textContent)).to.deep.equal({
+        mode: 'dark',
+        lightColorScheme: 'paper',
+        darkColorScheme: 'dim',
+        colorScheme: 'dim',
+      });
+    });
+
     it('reset colorScheme', () => {
       const Data = () => {
         const { setColorScheme, ...data } = useCurrentColorScheme({
@@ -411,9 +478,9 @@ describe('useCurrentColorScheme', () => {
 
       fireEvent.click(container.firstChild);
 
-      expect(global.localStorage.setItem.calledWith(DEFAULT_MODE_STORAGE_KEY, 'dark')).to.equal(
-        true,
-      );
+      expect(
+        global.localStorage.setItem.lastCall.calledWith(DEFAULT_MODE_STORAGE_KEY, 'dark'),
+      ).to.equal(true);
     });
 
     it('save system mode', () => {
@@ -624,5 +691,66 @@ describe('useCurrentColorScheme', () => {
         colorScheme: 'dark-dim',
       });
     });
+
+    it('reset mode in storage', () => {
+      const Data = () => {
+        const { setMode } = useCurrentColorScheme({
+          defaultMode: 'system',
+          defaultLightColorScheme: 'light',
+          defaultDarkColorScheme: 'dark',
+          supportedColorSchemes: ['light', 'dark'],
+        });
+        return (
+          <div>
+            <button data-testid="dark" onClick={() => setMode('dark')} />
+            <button data-testid="reset" onClick={() => setMode(null)} />
+          </div>
+        );
+      };
+      render(<Data />);
+
+      fireEvent.click(screen.getByTestId('dark'));
+
+      fireEvent.click(screen.getByTestId('reset'));
+
+      expect(
+        global.localStorage.setItem.lastCall.calledWith(DEFAULT_MODE_STORAGE_KEY, 'system'),
+      ).to.equal(true);
+    });
+
+    it('reset color scheme in storage', () => {
+      const Data = () => {
+        const { setColorScheme } = useCurrentColorScheme({
+          defaultMode: 'system',
+          defaultLightColorScheme: 'light',
+          defaultDarkColorScheme: 'dark',
+          supportedColorSchemes: ['light', 'dark'],
+        });
+        return (
+          <div>
+            <button data-testid="dark" onClick={() => setColorScheme('dark')} />
+            <button data-testid="reset" onClick={() => setColorScheme(null)} />
+          </div>
+        );
+      };
+      render(<Data />);
+
+      fireEvent.click(screen.getByTestId('dark'));
+
+      global.localStorage.setItem.resetHistory();
+      expect(global.localStorage.setItem.callCount).to.equal(0); // reset the calls to neglect inital setItem in the assertion below
+
+      fireEvent.click(screen.getByTestId('reset'));
+
+      expect(
+        global.localStorage.setItem.calledWith(
+          `${DEFAULT_COLOR_SCHEME_STORAGE_KEY}-light`,
+          'light',
+        ),
+      ).to.equal(true);
+      expect(
+        global.localStorage.setItem.calledWith(`${DEFAULT_COLOR_SCHEME_STORAGE_KEY}-dark`, 'dark'),
+      ).to.equal(true);
+    });
   });
 });
diff --git a/packages/mui-system/src/cssVars/useCurrentColorScheme.ts b/packages/mui-system/src/cssVars/useCurrentColorScheme.ts
@@ -90,13 +90,17 @@ export function getColorScheme<SupportedColorScheme extends string>(
   });
 }
 
-function resolveValue(key: string, defaultValue?: string) {
+function initializeValue(key: string, defaultValue: string) {
   if (typeof window === 'undefined') {
     return undefined;
   }
   let value;
   try {
     value = localStorage.getItem(key) || undefined;
+    if (!value) {
+      // the first time that user enters the site.
+      localStorage.setItem(key, defaultValue);
+    }
   } catch (e) {
     // Unsupported
   }
@@ -129,12 +133,20 @@ export default function useCurrentColorScheme<SupportedColorScheme extends strin
   const joinedColorSchemes = supportedColorSchemes.join(',');
 
   const [state, setState] = React.useState(() => {
-    const initialMode = resolveValue(modeStorageKey, defaultMode);
+    const initialMode = initializeValue(modeStorageKey, defaultMode);
+    const lightColorScheme = initializeValue(
+      `${colorSchemeStorageKey}-light`,
+      defaultLightColorScheme,
+    );
+    const darkColorScheme = initializeValue(
+      `${colorSchemeStorageKey}-dark`,
+      defaultDarkColorScheme,
+    );
     return {
       mode: initialMode,
       systemMode: getSystemMode(initialMode),
-      lightColorScheme: resolveValue(`${colorSchemeStorageKey}-light`) || defaultLightColorScheme,
-      darkColorScheme: resolveValue(`${colorSchemeStorageKey}-dark`) || defaultDarkColorScheme,
+      lightColorScheme,
+      darkColorScheme,
     } as State<SupportedColorScheme>;
   });
 
@@ -143,10 +155,11 @@ export default function useCurrentColorScheme<SupportedColorScheme extends strin
   const setMode: Result<SupportedColorScheme>['setMode'] = React.useCallback(
     (mode) => {
       setState((currentState) => {
-        const newMode = !mode ? defaultMode : mode;
         if (mode === currentState.mode) {
+          // do nothing if mode does not change
           return currentState;
         }
+        const newMode = !mode ? defaultMode : mode;
         try {
           localStorage.setItem(modeStorageKey, newMode);
         } catch (e) {
@@ -164,18 +177,26 @@ export default function useCurrentColorScheme<SupportedColorScheme extends strin
 
   const setColorScheme: Result<SupportedColorScheme>['setColorScheme'] = React.useCallback(
     (value) => {
-      if (!value || typeof value === 'string') {
+      if (!value) {
+        setState((currentState) => {
+          try {
+            localStorage.setItem(`${colorSchemeStorageKey}-light`, defaultLightColorScheme);
+            localStorage.setItem(`${colorSchemeStorageKey}-dark`, defaultDarkColorScheme);
+          } catch (e) {
+            // Unsupported
+          }
+          return {
+            ...currentState,
+            lightColorScheme: defaultLightColorScheme,
+            darkColorScheme: defaultDarkColorScheme,
+          };
+        });
+      } else if (typeof value === 'string') {
         if (value && !joinedColorSchemes.includes(value)) {
           console.error(`\`${value}\` does not exist in \`theme.colorSchemes\`.`);
         } else {
           setState((currentState) => {
             const newState = { ...currentState };
-            if (!value) {
-              // reset to default color scheme
-              newState.lightColorScheme = defaultLightColorScheme;
-              newState.darkColorScheme = defaultDarkColorScheme;
-              return newState;
-            }
             processState(currentState, (mode) => {
               try {
                 localStorage.setItem(`${colorSchemeStorageKey}-${mode}`, value);
@@ -192,33 +213,40 @@ export default function useCurrentColorScheme<SupportedColorScheme extends strin
             return newState;
           });
         }
-      } else if (
-        (value.light && !joinedColorSchemes.includes(value.light)) ||
-        (value.dark && !joinedColorSchemes.includes(value.dark))
-      ) {
-        console.error(`\`${value}\` does not exist in \`theme.colorSchemes\`.`);
       } else {
         setState((currentState) => {
           const newState = { ...currentState };
-          if (value.light || value.light === null) {
-            newState.lightColorScheme =
-              value.light === null ? defaultLightColorScheme : value.light;
+          const newLightColorScheme = value.light === null ? defaultLightColorScheme : value.light;
+          const newDarkColorScheme = value.dark === null ? defaultDarkColorScheme : value.dark;
+
+          if (newLightColorScheme) {
+            if (!joinedColorSchemes.includes(newLightColorScheme)) {
+              console.error(`\`${newLightColorScheme}\` does not exist in \`theme.colorSchemes\`.`);
+            } else {
+              newState.lightColorScheme = newLightColorScheme;
+              try {
+                localStorage.setItem(`${colorSchemeStorageKey}-light`, newLightColorScheme);
+              } catch (error) {
+                // Unsupported
+              }
+            }
           }
-          if (value.dark || value.dark === null) {
-            newState.darkColorScheme = value.dark === null ? defaultDarkColorScheme : value.dark;
+
+          if (newDarkColorScheme) {
+            if (!joinedColorSchemes.includes(newDarkColorScheme)) {
+              console.error(`\`${newDarkColorScheme}\` does not exist in \`theme.colorSchemes\`.`);
+            } else {
+              newState.darkColorScheme = newDarkColorScheme;
+              try {
+                localStorage.setItem(`${colorSchemeStorageKey}-dark`, newDarkColorScheme);
+              } catch (error) {
+                // Unsupported
+              }
+            }
           }
+
           return newState;
         });
-        try {
-          if (value.light) {
-            localStorage.setItem(`${colorSchemeStorageKey}-light`, value.light);
-          }
-          if (value.dark) {
-            localStorage.setItem(`${colorSchemeStorageKey}-dark`, value.dark);
-          }
-        } catch (e) {
-          // Unsupported
-        }
       }
     },
     [joinedColorSchemes, colorSchemeStorageKey, defaultLightColorScheme, defaultDarkColorScheme],
@@ -253,25 +281,6 @@ export default function useCurrentColorScheme<SupportedColorScheme extends strin
     return () => media.removeListener(handler);
   }, []);
 
-  // Save mode, lightColorScheme & darkColorScheme to localStorage
-  React.useEffect(() => {
-    try {
-      if (state.mode) {
-        localStorage.setItem(modeStorageKey, state.mode);
-      }
-      processState(state, (mode) => {
-        if (mode === 'light') {
-          localStorage.setItem(`${colorSchemeStorageKey}-light`, state.lightColorScheme);
-        }
-        if (mode === 'dark') {
-          localStorage.setItem(`${colorSchemeStorageKey}-dark`, state.darkColorScheme);
-        }
-      });
-    } catch (e) {
-      // Unsupported
-    }
-  }, [state, colorSchemeStorageKey, modeStorageKey]);
-
   // Handle when localStorage has changed
   React.useEffect(() => {
     const handleStorage = (event: StorageEvent) => {

diff --git a/test/regressions/fixtures/CssVarsProvider/ColorSchemeSelector.js b/test/regressions/fixtures/CssVarsProvider/ColorSchemeSelector.js
@@ -15,12 +15,13 @@ const { CssVarsProvider } = createCssVarsProvider({
   },
   defaultColorScheme: {
     light: 'light',
+    dark: 'dark',
   },
 });
 
-export default function DarkModeSpecificity() {
+export default function ColorSchemeSelector() {
   return (
-    <CssVarsProvider modeStorageKey="dark-mode-specificity">
+    <CssVarsProvider>
       <Box
         sx={(theme) => ({
           p: 2,

diff --git a/test/regressions/fixtures/CssVarsProvider/DarkModeSpecificity.js b/test/regressions/fixtures/CssVarsProvider/DarkModeSpecificity.js
@@ -33,7 +33,7 @@ const DarkMode = () => {
 
 export default function DarkModeSpecificity() {
   return (
-    <CssVarsProvider modeStorageKey="dark-mode-specificity">
+    <CssVarsProvider>
       <DarkMode />
       <div style={{ background: 'var(--background-default)', color: '#888', padding: '1rem' }}>
         Background should be #000.

diff --git a/test/regressions/fixtures/CssVarsProvider/ForceColorSchemes.js b/test/regressions/fixtures/CssVarsProvider/ForceColorSchemes.js
@@ -22,9 +22,9 @@ const { CssVarsProvider } = createCssVarsProvider({
   },
 });
 
-export default function DarkModeSpecificity() {
+export default function ForceColorSchemes() {
   return (
-    <CssVarsProvider modeStorageKey="force-color-schemes">
+    <CssVarsProvider>
       <div
         data-color-scheme="dark"
         style={{

diff --git a/test/regressions/index.test.js b/test/regressions/index.test.js
@@ -77,6 +77,12 @@ async function main() {
   await fse.emptyDir(screenshotDir);
 
   describe('visual regressions', () => {
+    beforeEach(async () => {
+      await page.evaluate(() => {
+        localStorage.clear();
+      });
+    });
+
     after(async () => {
       await browser.close();
     });