From f1bb3a6dfe986e25fc265fac7890d387d56f5b4e Mon Sep 17 00:00:00 2001
From: Josh C <bowheart31@gmail.com>
Date: Tue, 30 Jul 2024 14:53:59 -0400
Subject: [PATCH] fix(react): restore `useAtomSelector` mounted state when
 swapping inline refs (#108)

---
 packages/react/src/hooks/useAtomSelector.ts   |  2 +-
 .../react/test/units/useAtomSelector.test.tsx | 48 +++++++++++++++++++
 2 files changed, 49 insertions(+), 1 deletion(-)

diff --git a/packages/react/src/hooks/useAtomSelector.ts b/packages/react/src/hooks/useAtomSelector.ts
index 200fd92a..46b4ceed 100644
--- a/packages/react/src/hooks/useAtomSelector.ts
+++ b/packages/react/src/hooks/useAtomSelector.ts
@@ -68,7 +68,7 @@ export const useAtomSelector = <T, Args extends any[]>(
       selectorOrConfig as AtomSelectorOrConfig<any, any[]>,
       resolvedArgs
     )
-    ;(render as any).mounted = false
+    ;(render as any).mounted = true
   }
 
   const cache = isSwappingRefs
diff --git a/packages/react/test/units/useAtomSelector.test.tsx b/packages/react/test/units/useAtomSelector.test.tsx
index b776cbd3..e79aa16e 100644
--- a/packages/react/test/units/useAtomSelector.test.tsx
+++ b/packages/react/test/units/useAtomSelector.test.tsx
@@ -567,4 +567,52 @@ describe('useAtomSelector', () => {
     expect(renders).toBe(6) // 3 rerenders + 3 for strict mode
     expect(ecosystem._graph.nodes).toMatchSnapshot()
   })
+
+  test('inline selector stays subscribed after being swapped out', async () => {
+    jest.useFakeTimers()
+    const atom1 = atom('1', () => ({ val: 1 }))
+    const selector = ({ get }: AtomGetters) => get(atom1)
+
+    function Test() {
+      const { val } = useAtomSelector({
+        resultsComparator: (a, b) => a.val === b.val,
+        selector,
+      })
+
+      return (
+        <>
+          <div data-testid="text">{val}</div>
+        </>
+      )
+    }
+
+    const { findByTestId } = renderInEcosystem(<Test />, {
+      useStrictMode: true,
+    })
+
+    const div = await findByTestId('text')
+
+    expect(div.innerHTML).toBe('1')
+
+    act(() => {
+      ecosystem.getInstance(atom1).setState({ val: 2 })
+      jest.runAllTimers()
+    })
+
+    expect(div.innerHTML).toBe('2')
+
+    act(() => {
+      ecosystem.getInstance(atom1).setState({ val: 3 })
+      jest.runAllTimers()
+    })
+
+    expect(div.innerHTML).toBe('3')
+
+    act(() => {
+      ecosystem.getInstance(atom1).setState({ val: 4 })
+      jest.runAllTimers()
+    })
+
+    expect(div.innerHTML).toBe('4')
+  })
 })