From 5c0118941280bb265ca7813afb987f89c8c97a17 Mon Sep 17 00:00:00 2001
From: Teague Stockwell <71202372+teaguestockwell@users.noreply.github.com>
Date: Thu, 30 Dec 2021 06:31:24 -0800
Subject: [PATCH] fix: useMedia SSR hydration bug with defaultState (#2216)

* fix: useMedia SSR hydration bug with defaultState

Prevent a React hydration mismatch when a default value is provided
by not defaulting to window.matchMedia(query).matches.

* Refactor nested ifs
---
 docs/useMedia.md       |  2 ++
 src/useMedia.ts        | 26 ++++++++++++++++++++++----
 tests/useMedia.test.ts | 42 ++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 66 insertions(+), 4 deletions(-)
 create mode 100644 tests/useMedia.test.ts

diff --git a/docs/useMedia.md b/docs/useMedia.md
index 1c99c442ab..97ee040c50 100644
--- a/docs/useMedia.md
+++ b/docs/useMedia.md
@@ -25,3 +25,5 @@ useMedia(query: string, defaultState: boolean = false): boolean;
 ```
 
 The `defaultState` parameter is only used as a fallback for server side rendering.
+
+When server side rendering, it is important to set this parameter because without it the server's initial state will fallback to false, but the client will initialize to the result of the media query. When React hydrates the server render, it may not match the client's state. See the [React docs](https://reactjs.org/docs/react-dom.html#hydrate) for more on why this is can lead to costly bugs 🐛.
diff --git a/src/useMedia.ts b/src/useMedia.ts
index 1cfa737572..5c34d4146a 100644
--- a/src/useMedia.ts
+++ b/src/useMedia.ts
@@ -1,10 +1,28 @@
 import { useEffect, useState } from 'react';
 import { isBrowser } from './misc/util';
 
-const useMedia = (query: string, defaultState: boolean = false) => {
-  const [state, setState] = useState(
-    isBrowser ? () => window.matchMedia(query).matches : defaultState
-  );
+const getInitialState = (query: string, defaultState?: boolean) => {
+  // Prevent a React hydration mismatch when a default value is provided by not defaulting to window.matchMedia(query).matches.
+  if (defaultState !== undefined) {
+    return defaultState;
+  }
+
+  if (isBrowser) {
+    return window.matchMedia(query).matches;
+  }
+  
+  // A default value has not been provided, and you are rendering on the server, warn of a possible hydration mismatch when defaulting to false.
+  if (process.env.NODE_ENV !== 'production') {
+    console.warn(
+      '`useMedia` When server side rendering, defaultState should be defined to prevent a hydration mismatches.'
+    );
+  }
+
+  return false;
+};
+
+const useMedia = (query: string, defaultState?: boolean) => {
+  const [state, setState] = useState(getInitialState(query, defaultState));
 
   useEffect(() => {
     let mounted = true;
diff --git a/tests/useMedia.test.ts b/tests/useMedia.test.ts
new file mode 100644
index 0000000000..052d8bb009
--- /dev/null
+++ b/tests/useMedia.test.ts
@@ -0,0 +1,42 @@
+import { renderHook } from '@testing-library/react-hooks';
+import { renderHook as renderHookSSR } from '@testing-library/react-hooks/server';
+import { useMedia } from '../src';
+
+const createMockMediaMatcher = (matches: Record<string, boolean>) => (qs: string) => ({
+  matches: matches[qs] ?? false,
+  addListener: () => {},
+  removeListener: () => {},
+});
+
+describe('useMedia', () => {
+  beforeEach(() => {
+    window.matchMedia = createMockMediaMatcher({
+      '(min-width: 500px)': true,
+      '(min-width: 1000px)': false,
+    }) as any;
+  });
+  it('should return true if media query matches', () => {
+    const { result } = renderHook(() => useMedia('(min-width: 500px)'));
+    expect(result.current).toBe(true);
+  });
+  it('should return false if media query does not match', () => {
+    const { result } = renderHook(() => useMedia('(min-width: 1200px)'));
+    expect(result.current).toBe(false);
+  });
+  it('should return default state before hydration', () => {
+    const { result } = renderHookSSR(() => useMedia('(min-width: 500px)', false));
+    expect(result.current).toBe(false);
+  });
+  it('should return media query result after hydration', async () => {
+    const { result, hydrate } = renderHookSSR(() => useMedia('(min-width: 500px)', false));
+    expect(result.current).toBe(false);
+    hydrate();
+    expect(result.current).toBe(true);
+  });
+  it('should return media query result after hydration', async () => {
+    const { result, hydrate } = renderHookSSR(() => useMedia('(min-width: 1200px)', true));
+    expect(result.current).toBe(true);
+    hydrate();
+    expect(result.current).toBe(false);
+  });
+});