fix teardowns

Author: jquense

src/useCustomEffect.ts

@@ -1,4 +1,12 @@
-import { DependencyList, EffectCallback, useRef, useEffect } from 'react'
+import {
+  DependencyList,
+  EffectCallback,
+  useRef,
+  useEffect,
+  useDebugValue,
+} from 'react'
+import useWillUnmount from './useWillUnmount'
+import useMounted from './useMounted'
 
 export type EffectHook = (effect: EffectCallback, deps?: DependencyList) => void
 
@@ -9,6 +17,11 @@ export type CustomEffectOptions<TDeps> = {
   effectHook?: EffectHook
 }
 
+type CleanUp = {
+  (): void
+  cleanup?: ReturnType<EffectCallback>
+}
+
 /**
  * a useEffect() hook with customized depedency comparision
  *
@@ -40,20 +53,39 @@ function useCustomEffect<TDeps extends DependencyList = DependencyList>(
   dependencies: TDeps,
   isEqualOrOptions: IsEqual<TDeps> | CustomEffectOptions<TDeps>,
 ) {
+  const isMounted = useMounted()
   const { isEqual, effectHook = useEffect } =
     typeof isEqualOrOptions === 'function'
       ? { isEqual: isEqualOrOptions }
       : isEqualOrOptions
 
-  const depsRef = useRef<TDeps | undefined>()
+  const dependenciesRef = useRef<TDeps>()
+  dependenciesRef.current = dependencies
+
+  const cleanupRef = useRef<CleanUp | null>(null)
+
   effectHook(() => {
-    const prev = depsRef.current
-    depsRef.current = dependencies
+    // If the ref the is `null` it's either the first effect or the last effect
+    // ran and was cleared, meaning _this_ update should run, b/c the equality
+    // check failed on in the cleanup of the last effect.
+    if (cleanupRef.current === null) {
+      const cleanup = effect()
 
-    if (!prev || !isEqual(dependencies, prev)) {
-      return effect()
+      cleanupRef.current = () => {
+        if (isMounted() && isEqual(dependenciesRef.current!, dependencies)) {
+          return
+        }
+
+        cleanupRef.current = null
+        if (cleanup) cleanup()
+      }
+      cleanupRef.current.cleanup = cleanup
     }
+
+    return cleanupRef.current ?? undefined
   })
+
+  useDebugValue(effect)
 }
 
 export default useCustomEffect

test/useCustomEffect.test.tsx

@@ -40,7 +40,7 @@ describe('useCustomEffect', () => {
     const spy = jest.fn()
     const hookSpy = jest.fn().mockImplementation(useImmediateUpdateEffect)
 
-    const [, wrapper] = renderHook(
+    renderHook(
       ({ value }) => {
         useCustomEffect(spy, [value], {
           isEqual: (next, prev) => next[0].foo === prev[0].foo,
@@ -50,6 +50,7 @@ describe('useCustomEffect', () => {
       { value: { foo: true } },
     )
 
+    // the update and unmount hook setup
     expect(hookSpy).toHaveBeenCalledTimes(1)
     // not called b/c useImmediateUpdateEffect doesn't run on initial render
     expect(spy).toHaveBeenCalledTimes(0)

test/useMutationObserver.test.tsx

@@ -44,7 +44,7 @@ describe('useMutationObserver', () => {
     disconnentSpy?.mockRestore()
   })
 
-  it.only('should update config', async () => {
+  it('should update config', async () => {
     const teardown = jest.fn()
     const spy = jest.fn(() => teardown)
 
@@ -82,6 +82,7 @@ describe('useMutationObserver', () => {
       ],
       expect.anything(),
     )
+    expect(disconnentSpy).toBeCalledTimes(1)
 
     wrapper.unmount()