Merge pull request #432 from huchenme/pr/after-unmount

Feature: add addCleanup util function

Author: mpeyper

.all-contributorsrc

@@ -183,6 +183,17 @@
       "contributions": [
         "test"
       ]
+    },
+    {
+      "login": "huchenme",
+      "name": "Hu Chen",
+      "avatar_url": "https://avatars3.githubusercontent.com/u/2078389?v=4",
+      "profile": "https://huchen.dev/",
+      "contributions": [
+        "code",
+        "doc",
+        "example"
+      ]
     }
   ],
   "commitConvention": "none"

README.md

@@ -164,11 +164,13 @@ Thanks goes to these wonderful people ([emoji key](https://allcontributors.org/d
     <td align="center"><a href="https://github.com/102"><img src="https://avatars1.githubusercontent.com/u/5839225?v=4?s=100" width="100px;" alt=""/><br /><sub><b>Roman Gusev</b></sub></a><br /><a href="https://github.com/testing-library/react-hooks-testing-library/commits?author=102" title="Documentation">📖</a></td>
     <td align="center"><a href="https://github.com/hemlok"><img src="https://avatars2.githubusercontent.com/u/9043345?v=4?s=100" width="100px;" alt=""/><br /><sub><b>Adam Seckel</b></sub></a><br /><a href="https://github.com/testing-library/react-hooks-testing-library/commits?author=hemlok" title="Code">💻</a></td>
     <td align="center"><a href="https://keiya01.github.io/portfolio"><img src="https://avatars1.githubusercontent.com/u/34934510?v=4?s=100" width="100px;" alt=""/><br /><sub><b>keiya sasaki</b></sub></a><br /><a href="https://github.com/testing-library/react-hooks-testing-library/commits?author=keiya01" title="Tests">⚠️</a></td>
+    <td align="center"><a href="https://huchen.dev/"><img src="https://avatars3.githubusercontent.com/u/2078389?v=4?s=100" width="100px;" alt=""/><br /><sub><b>Hu Chen</b></sub></a><br /><a href="https://github.com/testing-library/react-hooks-testing-library/commits?author=huchenme" title="Code">💻</a> <a href="https://github.com/testing-library/react-hooks-testing-library/commits?author=huchenme" title="Documentation">📖</a> <a href="#example-huchenme" title="Examples">💡</a></td>
   </tr>
 </table>
 
-<!-- markdownlint-enable -->
+<!-- markdownlint-restore -->
 <!-- prettier-ignore-end -->
+
 <!-- ALL-CONTRIBUTORS-LIST:END -->
 
 This project follows the [all-contributors](https://allcontributors.org/) specification.

docs/api-reference.md

@@ -10,6 +10,7 @@ route: '/reference/api'
 - [`renderHook`](/reference/api#renderhook)
 - [`act`](/reference/api#act)
 - [`cleanup`](/reference/api#cleanup)
+- [`addCleanup`](/reference/api#addcleanup)
 
 ---
 
@@ -108,7 +109,9 @@ This is the same [`act` function](https://reactjs.org/docs/test-utils.html#act)
 function cleanup: Promise<void>
 ```
 
-Unmounts any rendered hooks rendered with `renderHook`, ensuring all effects have been flushed.
+Unmounts any rendered hooks rendered with `renderHook`, ensuring all effects have been flushed. Any
+callbacks added with [`addCleanup`](<(/reference/api#addCleanup).>) will also be called when
+`cleanup` is run.
 
 > Please note that this is done automatically if the testing framework you're using supports the
 > `afterEach` global (like Jest, mocha and Jasmine). If not, you will need to do manual cleanups
@@ -147,6 +150,35 @@ variable to `true` before importing `@testing-library/react-hooks` will also dis
 
 ---
 
+## `addCleanup`
+
+```js
+function addCleanup(callback: function(): void|Promise<void>): function(): void
+```
+
+Add a callback to be called during [`cleanup`](/reference/api#cleanup), returning a function to
+remove the cleanup if is no longer required. Cleanups are called in reverse order to being added.
+This is usually only relevant when wanting a cleanup to run after the component has been unmounted.
+
+If the provided callback is an `async` function or returns a promise, `cleanup` will wait for it to
+be resolved before moving onto the next cleanup callback.
+
+> Please note that any cleanups added using `addCleanup` are removed after `cleanup` is called. For
+> cleanups that need to run with every test, it is advised to add them in a `beforeEach` block (or
+> equivalent for your test runner).
+
+## `removeCleanup`
+
+```js
+function removeCleanup(callback: function(): void|Promise<void>): void
+```
+
+Removes a cleanup callback previously added with [`addCleanup`](/reference/api#addCleanup). Once
+removed, the provided callback will no longer execute as part of running
+[`cleanup`](/reference/api#cleanup).
+
+---
+
 ## Async Utilities
 
 ### `waitForNextUpdate`

src/cleanup.js

@@ -4,12 +4,15 @@ let cleanupCallbacks = []
 
 async function cleanup() {
   await flushMicroTasks()
-  cleanupCallbacks.forEach((cb) => cb())
+  for (const callback of cleanupCallbacks) {
+    await callback()
+  }
   cleanupCallbacks = []
 }
 
 function addCleanup(callback) {
-  cleanupCallbacks.push(callback)
+  cleanupCallbacks.unshift(callback)
+  return () => removeCleanup(callback)
 }
 
 function removeCleanup(callback) {

src/pure.js

@@ -99,4 +99,4 @@ function renderHook(callback, { initialProps, wrapper } = {}) {
   }
 }
 
-export { renderHook, cleanup, act }
+export { renderHook, cleanup, addCleanup, removeCleanup, act }

test/cleanup.test.js

@@ -1,5 +1,5 @@
 import { useEffect } from 'react'
-import { renderHook, cleanup } from 'src'
+import { renderHook, cleanup, addCleanup, removeCleanup } from 'src/pure'
 
 describe('cleanup tests', () => {
   test('should flush effects on cleanup', async () => {
@@ -38,4 +38,98 @@ describe('cleanup tests', () => {
     expect(cleanupCalled[1]).toBe(true)
     expect(cleanupCalled[2]).toBe(true)
   })
+
+  test('should call cleanups in reverse order', async () => {
+    let callSequence = []
+    addCleanup(() => {
+      callSequence.push('cleanup')
+    })
+    addCleanup(() => {
+      callSequence.push('another cleanup')
+    })
+    const hookWithCleanup = () => {
+      useEffect(() => {
+        return () => {
+          callSequence.push('unmount')
+        }
+      })
+    }
+    renderHook(() => hookWithCleanup())
+
+    await cleanup()
+
+    expect(callSequence).toEqual(['unmount', 'another cleanup', 'cleanup'])
+  })
+
+  test('should wait for async cleanup', async () => {
+    let callSequence = []
+    addCleanup(() => {
+      callSequence.push('cleanup')
+    })
+    addCleanup(async () => {
+      await new Promise((resolve) => setTimeout(resolve, 10))
+      callSequence.push('another cleanup')
+    })
+    const hookWithCleanup = () => {
+      useEffect(() => {
+        return () => {
+          callSequence.push('unmount')
+        }
+      })
+    }
+    renderHook(() => hookWithCleanup())
+
+    await cleanup()
+
+    expect(callSequence).toEqual(['unmount', 'another cleanup', 'cleanup'])
+  })
+
+  test('should remove cleanup using removeCleanup', async () => {
+    let callSequence = []
+    addCleanup(() => {
+      callSequence.push('cleanup')
+    })
+    const anotherCleanup = () => {
+      callSequence.push('another cleanup')
+    }
+    addCleanup(anotherCleanup)
+    const hookWithCleanup = () => {
+      useEffect(() => {
+        return () => {
+          callSequence.push('unmount')
+        }
+      })
+    }
+    renderHook(() => hookWithCleanup())
+
+    removeCleanup(anotherCleanup)
+
+    await cleanup()
+
+    expect(callSequence).toEqual(['unmount', 'cleanup'])
+  })
+
+  test('should remove cleanup using returned handler', async () => {
+    let callSequence = []
+    addCleanup(() => {
+      callSequence.push('cleanup')
+    })
+    const remove = addCleanup(() => {
+      callSequence.push('another cleanup')
+    })
+    const hookWithCleanup = () => {
+      useEffect(() => {
+        return () => {
+          callSequence.push('unmount')
+        }
+      })
+    }
+    renderHook(() => hookWithCleanup())
+
+    remove()
+
+    await cleanup()
+
+    expect(callSequence).toEqual(['unmount', 'cleanup'])
+  })
 })