feat(useImmediateUpdateEffect): allow teardown from last effect like builtin effect hooks

Author: jquense

src/useImmediateUpdateEffect.ts

@@ -1,5 +1,6 @@
-import { DependencyList, useRef } from 'react'
+import { DependencyList, useRef, EffectCallback } from 'react'
 import useStableMemo from './useStableMemo'
+import useWillUnmount from './useWillUnmount'
 
 /**
  * An _immediate_ effect that runs an effect callback when its dependency array
@@ -17,16 +18,25 @@ import useStableMemo from './useStableMemo'
  *
  * @category effects
  */
-function useImmediateUpdateEffect(effect: () => void, deps: DependencyList) {
+function useImmediateUpdateEffect(
+  effect: EffectCallback,
+  deps: DependencyList,
+) {
   const firstRef = useRef(true)
+  const tearDown = useRef<ReturnType<EffectCallback>>()
+
+  useWillUnmount(() => {
+    if (tearDown.current) tearDown.current()
+  })
 
   useStableMemo(() => {
     if (firstRef.current) {
       firstRef.current = false
       return
     }
 
-    effect()
+    if (tearDown.current) tearDown.current()
+    tearDown.current = effect()
   }, deps)
 }
 

src/useStableMemo.ts

@@ -1,4 +1,4 @@
-import { DependencyList, useEffect, useRef } from 'react'
+import { DependencyList, useRef } from 'react'
 
 function isEqual(a: DependencyList, b: DependencyList) {
   if (a.length !== b.length) return false

test/useCustomEffect.test.tsx

@@ -0,0 +1,57 @@
+import useCustomEffect from '../src/useCustomEffect'
+import useImmediateUpdateEffect from '../src/useImmediateUpdateEffect'
+import { renderHook } from './helpers'
+
+describe('useCustomEffect', () => {
+  it('should run custom isEqual logic', () => {
+    const teardown = jest.fn()
+
+    const spy = jest.fn().mockImplementation(() => teardown)
+    const isEqual = jest.fn((next, prev) => next[0].foo === prev[0].foo)
+
+    const [, wrapper] = renderHook(
+      ({ value }) => {
+        useCustomEffect(spy, [value], isEqual)
+      },
+      { value: { foo: true } },
+    )
+
+    expect(spy).toHaveBeenCalledTimes(1)
+
+    // matches isEqual
+    wrapper.setProps({ value: { foo: true } })
+
+    expect(spy).toHaveBeenCalledTimes(1)
+
+    // update that should trigger
+    wrapper.setProps({ value: { foo: false } })
+
+    expect(spy).toHaveBeenCalledTimes(2)
+    expect(isEqual).toHaveBeenCalledTimes(2)
+
+    expect(teardown).toBeCalledTimes(1)
+    expect(spy).toHaveBeenCalledTimes(2)
+
+    wrapper.unmount()
+    expect(teardown).toBeCalledTimes(2)
+  })
+
+  it('should accept different hooks', () => {
+    const spy = jest.fn()
+    const hookSpy = jest.fn().mockImplementation(useImmediateUpdateEffect)
+
+    const [, wrapper] = renderHook(
+      ({ value }) => {
+        useCustomEffect(spy, [value], {
+          isEqual: (next, prev) => next[0].foo === prev[0].foo,
+          effectHook: hookSpy,
+        })
+      },
+      { value: { foo: true } },
+    )
+
+    expect(hookSpy).toHaveBeenCalledTimes(1)
+    // not called b/c useImmediateUpdateEffect doesn't run on initial render
+    expect(spy).toHaveBeenCalledTimes(0)
+  })
+})

test/useImmediateUpdateEffect.test.tsx

@@ -3,7 +3,8 @@ import { renderHook } from './helpers'
 
 describe('useImmediateUpdateEffect', () => {
   it('should run update after value changes', () => {
-    const spy = jest.fn()
+    const teardown = jest.fn()
+    const spy = jest.fn().mockImplementation(() => teardown)
 
     const [, wrapper] = renderHook(
       ({ value }) => {
@@ -18,8 +19,18 @@ describe('useImmediateUpdateEffect', () => {
 
     expect(spy).toHaveBeenCalledTimes(1)
 
+    // update that doesn't change the deps Array
     wrapper.setProps({ value: 2, other: true })
 
     expect(spy).toHaveBeenCalledTimes(1)
+
+    // second update
+    wrapper.setProps({ value: 4, other: true })
+
+    expect(teardown).toBeCalledTimes(1)
+    expect(spy).toHaveBeenCalledTimes(2)
+
+    wrapper.unmount()
+    expect(teardown).toBeCalledTimes(2)
   })
 })