fix: #346 ensure setTheme Receives the Latest State (#347)

* fix: change setTheme function to correct handle passed callback
- Updated setTheme to properly handle passed callback and always pass relevant state value to callback
- Added saveToLS function to save values to localstorage

* test: add unit test to check the provided setTheme fix

Author: Junk-debug

next-themes/__tests__/index.test.tsx

@@ -459,6 +459,32 @@ describe('setTheme', () => {
     expect(result.current.theme).toBe('light')
     expect(result.current.resolvedTheme).toBe('light')
   })
+
+  test('setTheme(<function>) gets relevant state value', () => {
+    const consoleSpy = vi.spyOn(console, 'log').mockImplementation(() => { })
+
+    const { result } = renderHook(() => useTheme(), {
+      wrapper: ({ children }) => <ThemeProvider defaultTheme="light">{children}</ThemeProvider>
+    })
+
+    act(() => {
+      result.current.setTheme((theme) => {
+        console.log('1', theme)
+        return theme === 'dark' ? 'light' : 'dark'
+      })
+      result.current.setTheme((theme) => {
+        console.log('2', theme)
+        return theme === 'light' ? 'dark' : 'light'
+      })
+    })
+
+    expect(consoleSpy).toHaveBeenCalledWith('1', 'light')
+    expect(consoleSpy).toHaveBeenCalledWith('2', 'dark')
+    expect(result.current.theme).toBe('light')
+
+    consoleSpy.mockRestore()
+  })
+
 })
 
 describe('inline script', () => {

next-themes/src/index.tsx

@@ -8,7 +8,16 @@ const colorSchemes = ['light', 'dark']
 const MEDIA = '(prefers-color-scheme: dark)'
 const isServer = typeof window === 'undefined'
 const ThemeContext = React.createContext<UseThemeProps | undefined>(undefined)
-const defaultContext: UseThemeProps = { setTheme: _ => {}, themes: [] }
+const defaultContext: UseThemeProps = { setTheme: _ => { }, themes: [] }
+
+const saveToLS = (storageKey: string, value: string) => {
+  // Save to storage
+  try {
+    localStorage.setItem(storageKey, value)
+  } catch (e) {
+    // Unsupported
+  }
+}
 
 export const useTheme = () => React.useContext(ThemeContext) ?? defaultContext
 
@@ -79,20 +88,20 @@ const Theme = ({
     enable?.()
   }, [nonce])
 
-  const setTheme = React.useCallback(
-    value => {
-      const newTheme = typeof value === 'function' ? value(theme) : value
-      setThemeState(newTheme)
+  const setTheme = React.useCallback(value => {
+    if (typeof value === 'function') {
+      setThemeState(prevTheme => {
+        const newTheme = value(prevTheme)
 
-      // Save to storage
-      try {
-        localStorage.setItem(storageKey, newTheme)
-      } catch (e) {
-        // Unsupported
-      }
-    },
-    [theme]
-  )
+        saveToLS(storageKey, newTheme)
+
+        return newTheme
+      })
+    } else {
+      setThemeState(value)
+      saveToLS(storageKey, value)
+    }
+  }, [])
 
   const handleMediaQuery = React.useCallback(
     (e: MediaQueryListEvent | MediaQueryList) => {
@@ -234,7 +243,7 @@ const disableAnimation = (nonce?: string) => {
 
   return () => {
     // Force restyle
-    ;(() => window.getComputedStyle(document.body))()
+    ; (() => window.getComputedStyle(document.body))()
 
     // Wait for next tick before removing
     setTimeout(() => {