From 2aabd377daf94b4f3143dbb4e4c52cff8847ce43 Mon Sep 17 00:00:00 2001
From: ian <yihengz@outlook.com>
Date: Thu, 8 Aug 2024 16:23:16 +0800
Subject: [PATCH] feat: useTheme add onChange callback

---
 jest.config.js                                |  2 +-
 match-media-mock.js                           | 13 ++++++++++++
 .../hooks/src/useTheme/__test__/index.test.ts | 16 +-------------
 packages/hooks/src/useTheme/index.ts          | 21 +++++++++++++------
 4 files changed, 30 insertions(+), 22 deletions(-)
 create mode 100644 match-media-mock.js

diff --git a/jest.config.js b/jest.config.js
index 23c1d38339..824400839e 100644
--- a/jest.config.js
+++ b/jest.config.js
@@ -14,7 +14,7 @@ module.exports = {
   testPathIgnorePatterns: ['/.history/'],
   modulePathIgnorePatterns: ['<rootDir>/package.json'],
   resetMocks: false,
-  setupFiles: ['./jest.setup.js', 'jest-localstorage-mock'],
+  setupFiles: ['./jest.setup.js', 'jest-localstorage-mock', './match-media-mock.js'],
   setupFilesAfterEnv: ['@testing-library/jest-dom/extend-expect'],
   transform: {
     '^.+\\.tsx?$': ['ts-jest', { tsconfig: 'tsconfig.json' }],
diff --git a/match-media-mock.js b/match-media-mock.js
new file mode 100644
index 0000000000..099fa51778
--- /dev/null
+++ b/match-media-mock.js
@@ -0,0 +1,13 @@
+Object.defineProperty(window, 'matchMedia', {
+  writable: true,
+  value: jest.fn().mockImplementation((query) => ({
+    matches: false,
+    media: query,
+    onchange: null,
+    addListener: jest.fn(), // deprecated
+    removeListener: jest.fn(), // deprecated
+    addEventListener: jest.fn(),
+    removeEventListener: jest.fn(),
+    dispatchEvent: jest.fn(),
+  })),
+});
diff --git a/packages/hooks/src/useTheme/__test__/index.test.ts b/packages/hooks/src/useTheme/__test__/index.test.ts
index a06e76baab..3343269439 100644
--- a/packages/hooks/src/useTheme/__test__/index.test.ts
+++ b/packages/hooks/src/useTheme/__test__/index.test.ts
@@ -1,19 +1,5 @@
-Object.defineProperty(window, 'matchMedia', {
-  writable: true,
-  value: jest.fn().mockImplementation((query) => ({
-    matches: false,
-    media: query,
-    onchange: null,
-    addListener: jest.fn(), // deprecated
-    removeListener: jest.fn(), // deprecated
-    addEventListener: jest.fn(),
-    removeEventListener: jest.fn(),
-    dispatchEvent: jest.fn(),
-  })),
-});
-
 import { act, renderHook } from '@testing-library/react';
-import { useTheme } from '../index';
+import useTheme from '../index';
 
 describe('useTheme', () => {
   test('themeMode init', () => {
diff --git a/packages/hooks/src/useTheme/index.ts b/packages/hooks/src/useTheme/index.ts
index 3fecec96a2..a11e4f64d0 100644
--- a/packages/hooks/src/useTheme/index.ts
+++ b/packages/hooks/src/useTheme/index.ts
@@ -9,19 +9,27 @@ export enum ThemeMode {
 
 export type ThemeModeType = `${ThemeMode}`;
 
+export type ThemeType = 'light' | 'dark';
+
 const matchMedia = window.matchMedia('(prefers-color-scheme: dark)');
 
-function useCurrentTheme() {
-  const [theme, setTheme] = useState<'light' | 'dark'>(() => {
-    return matchMedia.matches ? ThemeMode.DARK : ThemeMode.LIGHT;
+type Callback = (theme: ThemeType) => void;
+
+function useCurrentTheme(callback: Callback = () => {}) {
+  const [theme, setTheme] = useState<ThemeType>(() => {
+    const init = matchMedia.matches ? ThemeMode.DARK : ThemeMode.LIGHT;
+    callback(init);
+    return init;
   });
 
   useEffect(() => {
     const onThemeChange: MediaQueryList['onchange'] = (event) => {
       if (event.matches) {
         setTheme(ThemeMode.DARK);
+        callback(ThemeMode.DARK);
       } else {
         setTheme(ThemeMode.LIGHT);
+        callback(ThemeMode.LIGHT);
       }
     };
 
@@ -30,17 +38,18 @@ function useCurrentTheme() {
     return () => {
       matchMedia.removeEventListener('change', onThemeChange);
     };
-  }, []);
+  }, [callback]);
 
   return theme;
 }
 
 type Options = {
   localStorageKey?: string;
+  onChange?: Callback;
 };
 
 export default function useTheme(options: Options = {}) {
-  const { localStorageKey } = options;
+  const { localStorageKey, onChange } = options;
 
   const [themeMode, setThemeMode] = useState<ThemeModeType>(() => {
     const preferredThemeMode =
@@ -57,7 +66,7 @@ export default function useTheme(options: Options = {}) {
     }
   };
 
-  const currentTheme = useCurrentTheme();
+  const currentTheme = useCurrentTheme(onChange);
   const theme = themeMode === ThemeMode.SYSTEM ? currentTheme : themeMode;
 
   return {