From da4bfddb26c680e2f6ea1bf712d1118320834e01 Mon Sep 17 00:00:00 2001
From: Tyler Swavely <tyler.swavely@nike.com>
Date: Tue, 19 Nov 2019 22:24:59 -0800
Subject: [PATCH] finish tests and hook

---
 src/useLocalStorage.ts        | 69 +++++++++++++++++++++--------------
 tests/useLocalStorage.test.ts | 45 ++++++++++++++++++-----
 2 files changed, 76 insertions(+), 38 deletions(-)

diff --git a/src/useLocalStorage.ts b/src/useLocalStorage.ts
index 051685be..c803c152 100644
--- a/src/useLocalStorage.ts
+++ b/src/useLocalStorage.ts
@@ -1,42 +1,55 @@
-import { useEffect, useState } from 'react';
 import { isClient } from './util';
+import { useMemo, useCallback, useEffect, Dispatch, SetStateAction } from 'react';
 
-type Dispatch<A> = (value: A) => void;
-type SetStateAction<S> = S | ((prevState: S) => S);
-
-const useLocalStorage = <T>(key: string, initialValue?: T, raw?: boolean): [T, Dispatch<SetStateAction<T>>] => {
-  if (!isClient) {
+const useLocalStorage = <T extends any>(
+  key: string,
+  initialValue?: any,
+  raw?: boolean
+): [any, Dispatch<SetStateAction<any>>] => {
+  if (!isClient || !localStorage) {
     return [initialValue as T, () => {}];
   }
 
-  const [state, setState] = useState<T>(() => {
+  let localStorageValue: string | null = null;
+  try {
+    localStorageValue = localStorage.getItem(key);
+  } catch {
+    // If user is in private mode or has storage restriction
+    // localStorage can throw.
+    localStorageValue = initialValue;
+  }
+
+  const state = useMemo(() => {
     try {
-      const localStorageValue = localStorage.getItem(key);
-      if (typeof localStorageValue !== 'string') {
-        localStorage.setItem(key, raw ? String(initialValue) : JSON.stringify(initialValue));
-        return initialValue;
-      } else {
-        return raw ? localStorageValue : JSON.parse(localStorageValue || 'null');
-      }
+      if (localStorageValue === null) return initialValue; // key hasn't been set yet
+      return raw ? localStorageValue : JSON.parse(localStorageValue);
     } catch {
-      // If user is in private mode or has storage restriction
-      // localStorage can throw. JSON.parse and JSON.stringify
-      // can throw, too.
-      return initialValue;
+      /* JSON.parse and JSON.stringify can throw. */
+      return localStorageValue === null ? initialValue : localStorageValue;
     }
-  });
+  }, [key, localStorageValue, initialValue]);
+
+  const setState = useCallback(
+    (valOrFunc: any) => {
+      try {
+        let newState = typeof valOrFunc === 'function' ? valOrFunc(state) : valOrFunc;
+        newState = typeof newState === 'string' ? newState : JSON.stringify(newState);
+        localStorage.setItem(key, newState);
+      } catch {
+        /**
+         * If user is in private mode or has storage restriction
+         * localStorage can throw. Also JSON.stringify can throw.
+         */
+      }
+    },
+    [state, raw]
+  );
 
   useEffect(() => {
-    try {
-      const serializedState = raw ? String(state) : JSON.stringify(state);
-      localStorage.setItem(key, serializedState);
-    } catch {
-      // If user is in private mode or has storage restriction
-      // localStorage can throw. Also JSON.stringify can throw.
-    }
-  }, [state]);
+    if (localStorageValue === null) setState(initialValue);
+  }, [localStorageValue, setState]);
 
-  return [state, setState];
+  return [state as any, setState];
 };
 
 export default useLocalStorage;
diff --git a/tests/useLocalStorage.test.ts b/tests/useLocalStorage.test.ts
index c4b23295..816cfa7e 100644
--- a/tests/useLocalStorage.test.ts
+++ b/tests/useLocalStorage.test.ts
@@ -55,7 +55,8 @@ describe(useLocalStorage, () => {
     rerender();
 
     const [foo] = result.current;
-    expect(foo).toEqual("baz");
+    expect(foo).not.toMatch(/\\/i); // should not contain extra escapes
+    expect(foo).toBe('baz');
   });
   it("keeps multiple hooks accessing the same key in sync", () => {
     localStorage.setItem("foo", "bar");
@@ -112,19 +113,43 @@ describe(useLocalStorage, () => {
     );
 
     const [, setFoo] = result.current;
-    act(() =>
-      setFoo(state => {
-        console.log(state);
-        return { ...state, fizz: "buzz" };
-      })
-    );
+    act(() => setFoo(state => ({ ...state, fizz: "buzz" })));
     rerender();
 
     const [value] = result.current;
-
-    console.log(value);
-
     expect(value.foo).toEqual("bar");
     expect(value.fizz).toEqual("buzz");
   });
+  describe("raw setting", () => {
+    it('returns a string when localStorage is a stringified object', () => {
+      localStorage.setItem('foo', JSON.stringify({ fizz: 'buzz' }));
+      const { result } = renderHook(() => useLocalStorage('foo', null, true));
+      const [foo] = result.current;
+      expect(typeof foo).toBe('string');
+    });
+    it('returns a string after an update', () => {
+      localStorage.setItem('foo', JSON.stringify({ fizz: 'buzz' }));
+      const { result, rerender } = renderHook(() => useLocalStorage('foo', null, true));
+
+      const [,setFoo] = result.current;
+      act(() => setFoo({ fizz: 'bang' }))
+      rerender();
+
+      const [foo] = result.current;
+      expect(typeof foo).toBe('string');
+      expect(JSON.parse(foo)).toBeInstanceOf(Object);
+      expect(JSON.parse(foo).fizz).toEqual('bang');
+    });
+    it('still forces setState to a string', () => {
+      localStorage.setItem('foo', JSON.stringify({ fizz: 'buzz' }));
+      const { result, rerender } = renderHook(() => useLocalStorage('foo', null, true));
+
+      const [,setFoo] = result.current;
+      act(() => setFoo({ fizz: 'bang' }))
+      rerender();
+
+      const [value] = result.current;
+      expect(JSON.parse(value).fizz).toEqual('bang');
+    });
+  });
 });