From 233ceb5356583fc22a7cd7745b1ddd8c53122159 Mon Sep 17 00:00:00 2001
From: Alf Eaton <alf.eaton@overleaf.com>
Date: Thu, 23 Sep 2021 11:35:50 +0100
Subject: [PATCH] Allow function as value for usePersistedState hook (#5131)

* Allow function value in usePersistedState
* Add tests for usePersistedState
* Use nullish coalescing to avoid calling getItem twice

GitOrigin-RevId: e0351addea904aefb7a402bff32689792b49fbbb
---
 .../js/shared/hooks/use-persisted-state.js    |  21 ++-
 .../shared/hooks/use-persisted-state.test.js  | 141 ++++++++++++++++++
 2 files changed, 154 insertions(+), 8 deletions(-)
 create mode 100644 services/web/test/frontend/shared/hooks/use-persisted-state.test.js

diff --git a/services/web/frontend/js/shared/hooks/use-persisted-state.js b/services/web/frontend/js/shared/hooks/use-persisted-state.js
index 3cc8f95e29..e009b013c7 100644
--- a/services/web/frontend/js/shared/hooks/use-persisted-state.js
+++ b/services/web/frontend/js/shared/hooks/use-persisted-state.js
@@ -3,18 +3,23 @@ import localStorage from '../../infrastructure/local-storage'
 
 function usePersistedState(key, defaultValue) {
   const [value, setValue] = useState(() => {
-    const keyExists = localStorage.getItem(key) != null
-    return keyExists ? localStorage.getItem(key) : defaultValue
+    return localStorage.getItem(key) ?? defaultValue
   })
 
   const updateFunction = useCallback(
     newValue => {
-      if (newValue === defaultValue) {
-        localStorage.removeItem(key)
-      } else {
-        localStorage.setItem(key, newValue)
-      }
-      setValue(newValue)
+      setValue(value => {
+        const actualNewValue =
+          typeof newValue === 'function' ? newValue(value) : newValue
+
+        if (actualNewValue === defaultValue) {
+          localStorage.removeItem(key)
+        } else {
+          localStorage.setItem(key, actualNewValue)
+        }
+
+        return actualNewValue
+      })
     },
     [key, defaultValue]
   )
diff --git a/services/web/test/frontend/shared/hooks/use-persisted-state.test.js b/services/web/test/frontend/shared/hooks/use-persisted-state.test.js
new file mode 100644
index 0000000000..5dcc4077ba
--- /dev/null
+++ b/services/web/test/frontend/shared/hooks/use-persisted-state.test.js
@@ -0,0 +1,141 @@
+import sinon from 'sinon'
+import { expect } from 'chai'
+import { useEffect } from 'react'
+import { render, screen } from '@testing-library/react'
+import usePersistedState from '../../../../frontend/js/shared/hooks/use-persisted-state'
+import localStorage from '../../../../frontend/js/infrastructure/local-storage'
+
+describe('usePersistedState', function () {
+  beforeEach(function () {
+    sinon.spy(global.localStorage, 'getItem')
+    sinon.spy(global.localStorage, 'removeItem')
+    sinon.spy(global.localStorage, 'setItem')
+  })
+
+  afterEach(function () {
+    sinon.restore()
+  })
+
+  it('reads the value from localStorage', function () {
+    const key = 'test'
+    localStorage.setItem(key, 'foo')
+    expect(global.localStorage.setItem).to.have.callCount(1)
+
+    const Test = () => {
+      const [value] = usePersistedState(key)
+
+      return <div>{value}</div>
+    }
+
+    render(<Test />)
+    screen.getByText('foo')
+
+    expect(global.localStorage.getItem).to.have.callCount(1)
+    expect(global.localStorage.removeItem).to.have.callCount(0)
+    expect(global.localStorage.setItem).to.have.callCount(1)
+
+    expect(localStorage.getItem(key)).to.equal('foo')
+  })
+
+  it('uses the default value without storing anything', function () {
+    const key = 'test:default'
+
+    const Test = () => {
+      const [value] = usePersistedState(key, 'foo')
+
+      return <div>{value}</div>
+    }
+
+    render(<Test />)
+    screen.getByText('foo')
+
+    expect(global.localStorage.getItem).to.have.callCount(1)
+    expect(global.localStorage.removeItem).to.have.callCount(0)
+    expect(global.localStorage.setItem).to.have.callCount(0)
+
+    expect(localStorage.getItem(key)).to.be.null
+  })
+
+  it('stores the new value in localStorage', function () {
+    const key = 'test:store'
+    localStorage.setItem(key, 'foo')
+    expect(global.localStorage.setItem).to.have.callCount(1)
+
+    const Test = () => {
+      const [value, setValue] = usePersistedState(key, 'bar')
+
+      useEffect(() => {
+        setValue('baz')
+      }, [setValue])
+
+      return <div>{value}</div>
+    }
+
+    render(<Test />)
+
+    screen.getByText('baz')
+
+    expect(global.localStorage.getItem).to.have.callCount(1)
+    expect(global.localStorage.removeItem).to.have.callCount(0)
+    expect(global.localStorage.setItem).to.have.callCount(2)
+
+    expect(localStorage.getItem(key)).to.equal('baz')
+  })
+
+  it('removes the value from localStorage if it equals the default value', function () {
+    const key = 'test:store-default'
+    localStorage.setItem(key, 'foo')
+    expect(global.localStorage.setItem).to.have.callCount(1)
+
+    const Test = () => {
+      const [value, setValue] = usePersistedState(key, 'bar')
+
+      useEffect(() => {
+        // set a different value
+        setValue('baz')
+        expect(localStorage.getItem(key)).to.equal('baz')
+
+        // set the default value again
+        setValue('bar')
+      }, [setValue])
+
+      return <div>{value}</div>
+    }
+
+    render(<Test />)
+
+    screen.getByText('bar')
+
+    expect(global.localStorage.getItem).to.have.callCount(2)
+    expect(global.localStorage.removeItem).to.have.callCount(1)
+    expect(global.localStorage.setItem).to.have.callCount(2)
+
+    expect(localStorage.getItem(key)).to.be.null
+  })
+
+  it('handles function values', function () {
+    const key = 'test:store'
+    localStorage.setItem(key, 'foo')
+    expect(global.localStorage.setItem).to.have.callCount(1)
+
+    const Test = () => {
+      const [value, setValue] = usePersistedState(key)
+
+      useEffect(() => {
+        setValue(value => value + 'bar')
+      }, [setValue])
+
+      return <div>{value}</div>
+    }
+
+    render(<Test />)
+
+    screen.getByText('foobar')
+
+    expect(global.localStorage.getItem).to.have.callCount(1)
+    expect(global.localStorage.removeItem).to.have.callCount(0)
+    expect(global.localStorage.setItem).to.have.callCount(2)
+
+    expect(localStorage.getItem(key)).to.equal('foobar')
+  })
+})