fix(utils/atomWithStorage): async storage support (#569)

* fix(utils/atomWithStorage): async storage support

* add tests

* expose createJSONStorage

* update size snapshot

.size-snapshot.json

@@ -14,16 +14,16 @@
     }
   },
   "utils.js": {
-    "bundled": 13354,
-    "minified": 6528,
-    "gzipped": 2437,
+    "bundled": 13604,
+    "minified": 6640,
+    "gzipped": 2497,
     "treeshaked": {
       "rollup": {
         "code": 28,
         "import_statements": 28
       },
       "webpack": {
-        "code": 1280
+        "code": 1305
       }
     }
   },

src/utils.ts

@@ -12,4 +12,8 @@ export { freezeAtom, freezeAtomCreator } from './utils/freezeAtom'
 export { splitAtom } from './utils/splitAtom'
 export { atomWithDefault } from './utils/atomWithDefault'
 export { waitForAll } from './utils/waitForAll'
-export { atomWithStorage, atomWithHash } from './utils/atomWithStorage'
+export {
+  atomWithStorage,
+  atomWithHash,
+  createJSONStorage,
+} from './utils/atomWithStorage'

src/utils/atomWithStorage.ts

@@ -10,18 +10,27 @@ type Storage<Value> = {
   subscribe?: (key: string, callback: (value: Value) => void) => Unsubscribe
 }
 
-const defaultStorage: Storage<unknown> = {
+type StringStorage = {
+  getItem: (key: string) => string | null | Promise<string | null>
+  setItem: (key: string, newValue: string) => void | Promise<void>
+}
+
+export const createJSONStorage = (
+  getStringStorage: () => StringStorage
+): Storage<unknown> => ({
   getItem: (key) => {
-    const storedValue = localStorage.getItem(key)
-    if (storedValue === null) {
-      throw new Error('no value stored')
+    const value = getStringStorage().getItem(key)
+    if (value instanceof Promise) {
+      return value.then((v) => JSON.parse(v || ''))
     }
-    return JSON.parse(storedValue)
+    return JSON.parse(value || '')
   },
   setItem: (key, newValue) => {
-    localStorage.setItem(key, JSON.stringify(newValue))
+    getStringStorage().setItem(key, JSON.stringify(newValue))
   },
-}
+})
+
+const defaultStorage = createJSONStorage(() => localStorage)
 
 export function atomWithStorage<Value>(
   key: string,
@@ -30,7 +39,11 @@ export function atomWithStorage<Value>(
 ): PrimitiveAtom<Value> {
   const getInitialValue = () => {
     try {
-      return storage.getItem(key)
+      const value = storage.getItem(key)
+      if (value instanceof Promise) {
+        return value.catch(() => initialValue)
+      }
+      return value
     } catch {
       return initialValue
     }

tests/utils/atomWithStorage.test.tsx

@@ -6,7 +6,7 @@ import { getTestProvider } from '../testUtils'
 
 const Provider = getTestProvider()
 
-describe('atomWithStorage', () => {
+describe('atomWithStorage (sync)', () => {
   const storageData: Record<string, number> = {
     count: 10,
   }
@@ -47,7 +47,9 @@ describe('atomWithStorage', () => {
     await findByText('count: 11')
     expect(storageData.count).toBe(11)
   })
+})
 
+describe('atomWithStorage (async)', () => {
   const asyncStorageData: Record<string, number> = {
     count: 10,
   }
@@ -86,11 +88,73 @@ describe('atomWithStorage', () => {
       </Provider>
     )
 
+    await findByText('loading')
     await findByText('count: 10')
 
     fireEvent.click(getByText('button'))
     await findByText('count: 11')
-    expect(storageData.count).toBe(11)
+    await new Promise((r) => setTimeout(r, 20))
+    expect(asyncStorageData.count).toBe(11)
+  })
+
+  it('async new count', async () => {
+    const countAtom = atomWithStorage('count2', 20, asyncDummyStorage)
+
+    const Counter: React.FC = () => {
+      const [count, setCount] = useAtom(countAtom)
+      return (
+        <>
+          <div>count: {count}</div>
+          <button onClick={() => setCount((c) => c + 1)}>button</button>
+        </>
+      )
+    }
+
+    const { findByText, getByText } = render(
+      <Provider>
+        <Suspense fallback="loading">
+          <Counter />
+        </Suspense>
+      </Provider>
+    )
+
+    await findByText('loading')
+    await findByText('count: 20')
+
+    fireEvent.click(getByText('button'))
+    await findByText('count: 21')
+    await new Promise((r) => setTimeout(r, 20))
+    expect(asyncStorageData.count2).toBe(21)
+  })
+
+  it('async new count with delayInit', async () => {
+    const countAtom = atomWithStorage('count3', 30, {
+      ...asyncDummyStorage,
+      delayInit: true,
+    })
+
+    const Counter: React.FC = () => {
+      const [count, setCount] = useAtom(countAtom)
+      return (
+        <>
+          <div>count: {count}</div>
+          <button onClick={() => setCount((c) => c + 1)}>button</button>
+        </>
+      )
+    }
+
+    const { findByText, getByText } = render(
+      <Provider>
+        <Counter />
+      </Provider>
+    )
+
+    await findByText('count: 30')
+
+    fireEvent.click(getByText('button'))
+    await findByText('count: 31')
+    await new Promise((r) => setTimeout(r, 20))
+    expect(asyncStorageData.count3).toBe(31)
   })
 })