feat(theme-provider): theme provider changes #499

src/tedi/providers/theme-provider/theme-provider.spec.tsx

@@ -18,6 +18,7 @@ describe('ThemeProvider', () => {
 
   const TestComponent = () => {
     const { theme, setTheme } = useTheme();
+
     return (
       <div>
         <p data-testid="theme">{theme}</p>
@@ -53,7 +54,7 @@ describe('ThemeProvider', () => {
     expect(document.cookie).toContain('tedi-theme=default');
   });
 
-  it('ignores existing localStorage theme and uses provided prop theme instead', () => {
+  it('uses theme from localStorage over provided prop theme', () => {
     localStorage.setItem('tedi-theme', 'dark');
 
     render(
@@ -62,11 +63,10 @@ describe('ThemeProvider', () => {
       </ThemeProvider>
     );
 
-    // Provider should still use "default" because prop theme overrides storage
-    expect(screen.getByTestId('theme')).toHaveTextContent('default');
-    expect(document.documentElement).toHaveClass('tedi-theme--default');
-    expect(localStorage.getItem('tedi-theme')).toBe('default');
-    expect(document.cookie).toContain('tedi-theme=default');
+    expect(screen.getByTestId('theme')).toHaveTextContent('dark');
+    expect(document.documentElement).toHaveClass('tedi-theme--dark');
+    expect(localStorage.getItem('tedi-theme')).toBe('dark');
+    expect(document.cookie).toContain('tedi-theme=dark');
   });
 
   it('updates theme and document classes when setTheme is called', () => {
@@ -76,10 +76,8 @@ describe('ThemeProvider', () => {
       </ThemeProvider>
     );
 
-    const button = screen.getByText('Set Dark');
-
     act(() => {
-      button.click();
+      screen.getByText('Set Dark').click();
     });
 
     expect(screen.getByTestId('theme')).toHaveTextContent('dark');
@@ -133,29 +131,31 @@ describe('ThemeProvider', () => {
     expect(document.cookie).toContain('tedi-theme=rit');
   });
 
-  it('does not set theme if invalid theme is provided', () => {
-    const TestComponentWithInvalidTheme = () => {
+  it('allows custom theme values', () => {
+    const TestComponentWithCustomTheme = () => {
       const { theme, setTheme } = useTheme();
+
       return (
         <div>
           <p data-testid="theme">{theme}</p>
-          <button onClick={() => setTheme('invalid' as Theme)}>Set Invalid</button>
+          <button onClick={() => setTheme('invalid' as Theme)}>Set Custom</button>
         </div>
       );
     };
 
     render(
       <ThemeProvider theme="default">
-        <TestComponentWithInvalidTheme />
+        <TestComponentWithCustomTheme />
       </ThemeProvider>
     );
 
     act(() => {
-      screen.getByText('Set Invalid').click();
+      screen.getByText('Set Custom').click();
     });
 
-    expect(screen.getByTestId('theme')).toHaveTextContent('default');
-    expect(document.documentElement).toHaveClass('tedi-theme--default');
-    expect(localStorage.getItem('tedi-theme')).toBe('default');
+    expect(screen.getByTestId('theme')).toHaveTextContent('invalid');
+    expect(document.documentElement).toHaveClass('tedi-theme--invalid');
+    expect(localStorage.getItem('tedi-theme')).toBe('invalid');
+    expect(document.cookie).toContain('tedi-theme=invalid');
   });
 });

src/tedi/providers/theme-provider/theme-provider.tsx

@@ -1,10 +1,10 @@
 'use client';
 
-import { createContext, useCallback, useContext, useEffect, useState } from 'react';
+import { createContext, useContext, useEffect, useState } from 'react';
 
-export type Theme = 'default' | 'dark' | 'rit' | 'muis';
+type TEDITheme = 'default' | 'dark';
+export type Theme = TEDITheme | string;
 
-const AVAILABLE_THEMES: Theme[] = ['default', 'dark', 'rit', 'muis'];
 const THEME_CLASS_PREFIX = 'tedi-theme--';
 const STORAGE_KEY = 'tedi-theme';
 
@@ -20,40 +20,44 @@ const ThemeContext = createContext<ThemeContextValue>({
 
 export const useTheme = () => useContext(ThemeContext);
 
-export const ThemeProvider = ({ theme: initialTheme, children }: { theme?: Theme; children: React.ReactNode }) => {
-  const [theme, setThemeState] = useState<Theme>(() => {
-    if (initialTheme !== undefined) {
-      return initialTheme;
-    }
+function getInitialTheme(initialTheme?: Theme): Theme {
+  if (typeof window === 'undefined') {
+    return initialTheme ?? 'default';
+  }
 
-    if (typeof window !== 'undefined') {
-      const saved = localStorage.getItem(STORAGE_KEY) as Theme | null;
-      if (saved && AVAILABLE_THEMES.includes(saved)) {
-        return saved;
-      }
-    }
+  const fromStorage = localStorage.getItem(STORAGE_KEY);
+  if (fromStorage) return fromStorage;
+
+  const cookieMatch = document.cookie
+    .split('; ')
+    .find((c) => c.startsWith(`${STORAGE_KEY}=`))
+    ?.split('=')[1];
+
+  if (cookieMatch) return cookieMatch;
+
+  return initialTheme ?? 'default';
+}
 
-    return 'default';
-  });
+export const ThemeProvider = ({ theme: initialTheme, children }: { theme?: Theme; children: React.ReactNode }) => {
+  const [theme, setTheme] = useState<Theme>(() => getInitialTheme(initialTheme));
 
   useEffect(() => {
     if (typeof document === 'undefined') return;
 
     const root = document.documentElement;
 
-    AVAILABLE_THEMES.forEach((t) => {
-      root.classList.toggle(`${THEME_CLASS_PREFIX}${t}`, t === theme);
-    });
+    for (let i = root.classList.length - 1; i >= 0; i--) {
+      const cls = root.classList.item(i);
+      if (cls?.startsWith(THEME_CLASS_PREFIX)) {
+        root.classList.remove(cls);
+      }
+    }
+
+    root.classList.add(`${THEME_CLASS_PREFIX}${theme}`);
 
     localStorage.setItem(STORAGE_KEY, theme);
-    document.cookie = `${STORAGE_KEY}=${theme};path=/;max-age=31536000`;
+    document.cookie = `${STORAGE_KEY}=${theme}; path=/; max-age=31536000; SameSite=Lax`;
   }, [theme]);
 
-  const setTheme = useCallback((newTheme: Theme) => {
-    if (AVAILABLE_THEMES.includes(newTheme)) {
-      setThemeState(newTheme);
-    }
-  }, []);
-
   return <ThemeContext.Provider value={{ theme, setTheme }}>{children}</ThemeContext.Provider>;
 };