feat: make the storage key reactive (#42)

* make the storage key reactive

* use a ref to minimize dependencies of callbacks
(will prevent unneccessary mounting/unmounting of storage)

* remove unnecessary import

* flesh out test for reactive key
- tests if storage is properly mounted and unmounted if the key changes

---------

Co-authored-by: L <[email protected]>

src/hook.ts

@@ -35,9 +35,37 @@ export const useStorage = <T = any>(rawKey: RawKey, onInit?: Setter<T>) => {
   // Use to ensure we don't set render state after unmounted
   const isMounted = useRef(false)
 
+  // Ref that stores the render state, in order to minimize dependencies of callbacks below
+  const renderValueRef = useRef(onInit instanceof Function ? onInit() : onInit)
+  useEffect(() => {
+    renderValueRef.current = renderValue
+  }, [renderValue])
+
   // Storage state
   const storageRef = useRef(isObjectKey ? rawKey.instance : new Storage())
 
+  // Save the value OR current rendering value into chrome storage
+  const setStoreValue = useCallback(
+    (v?: T) =>
+      storageRef.current.set(key, v !== undefined ? v : renderValueRef.current),
+    [key]
+  )
+
+  // Store the value into chrome storage, then set its render state
+  const persistValue = useCallback(
+    async (setter: Setter<T>) => {
+      const newValue =
+        setter instanceof Function ? setter(renderValueRef.current) : setter
+
+      await setStoreValue(newValue)
+
+      if (isMounted.current) {
+        setRenderValue(newValue)
+      }
+    },
+    [setStoreValue]
+  )
+
   useEffect(() => {
     isMounted.current = true
     const watchConfig: StorageCallbackMap = {
@@ -65,32 +93,12 @@ export const useStorage = <T = any>(rawKey: RawKey, onInit?: Setter<T>) => {
       isMounted.current = false
       storageRef.current.unwatch(watchConfig)
     }
-  }, [])
-
-  // Save the value OR current rendering value into chrome storage
-  const setStoreValue = useCallback(
-    (v?: T) => storageRef.current.set(key, v !== undefined ? v : renderValue),
-    [renderValue]
-  )
-
-  // Store the value into chrome storage, then set its render state
-  const persistValue = useCallback(
-    async (setter: Setter<T>) => {
-      const newValue = setter instanceof Function ? setter(renderValue) : setter
-
-      await setStoreValue(newValue)
-
-      if (isMounted.current) {
-        setRenderValue(newValue)
-      }
-    },
-    [renderValue, setStoreValue]
-  )
+  }, [key, persistValue])
 
   const remove = useCallback(() => {
     storageRef.current.remove(key)
     setRenderValue(undefined)
-  }, [setRenderValue])
+  }, [key])
 
   return [
     renderValue,

src/index.test.ts

@@ -4,7 +4,7 @@
  * This module share storage between chrome storage and local storage.
  */
 import { beforeEach, describe, expect, jest, test } from "@jest/globals"
-import { act, renderHook } from "@testing-library/react"
+import { act, renderHook, waitFor } from "@testing-library/react"
 
 import type { StorageWatchEventListener } from "~index"
 
@@ -182,6 +182,55 @@ describe("react hook", () => {
 
     expect(removeListener).toHaveBeenCalled()
   })
+
+  test("is reactive to key changes", async () => {
+    const { setTriggers, getTriggers } = createStorageMock()
+
+    const key1 = "key1"
+    const key2 = "key2"
+    const initValue = "hello"
+    const key1Value = "hello world"
+    const key2Value = "hello world 2"
+
+    const { result, rerender, unmount } = renderHook(
+      ({ key }) => useStorage(key, initValue),
+      {
+        initialProps: { key: key1 }
+      }
+    )
+
+    // with initial key, set new value
+    await act(async () => {
+      await result.current[1](key1Value)
+    })
+    expect(setTriggers).toHaveBeenCalledWith({
+      key1: JSON.stringify(key1Value)
+    })
+
+    // re-render with new key, and ensure new key is looked up from storage and that we reset to initial value
+    await act(async () => {
+      rerender({ key: key2 })
+    })
+    expect(getTriggers).toHaveBeenCalledWith(key2)
+    await waitFor(() => expect(result.current[0]).toBe(initValue))
+
+    // set new key to new value
+    await act(async () => {
+      await result.current[1](key2Value)
+    })
+    expect(setTriggers).toHaveBeenCalledWith({
+      key2: JSON.stringify(key2Value)
+    })
+    await waitFor(() => expect(result.current[0]).toBe(key2Value))
+
+    // re-render with old key, and ensure old key's up-to-date value is fetched
+    await act(async () => {
+      rerender({ key: key1 })
+    })
+    await waitFor(() => expect(result.current[0]).toBe(key1Value))
+
+    unmount()
+  })
 })
 
 describe("watch/unwatch", () => {