Refactor addCleanup to not treat unmount differently for external cleanups

Author: mpeyper

docs/api-reference.md

@@ -109,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
@@ -151,26 +153,30 @@ variable to `true` before importing `@testing-library/react-hooks` will also dis
 ## `addCleanup`
 
 ```js
-function addCleanup(
-  callback: function(props?: any): any
-): void
+function addCleanup(callback: function(): void|Promise<void>): function(): void
 ```
 
-Callback to be called after `cleanup`.
+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.
 
-In some cases you might want to run some callback after internal `cleanup` happen, especially after
-`unmount` happens in `cleanup`. If the sequence matters to you, you could use `addCleanup`.
+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.
 
-```js
-import { addCleanup } from '@testing-library/react-hooks'
+> 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).
 
-jest.useFakeTimers()
+## `removeCleanup`
 
-addCleanup(() => {
-  jest.runOnlyPendingTimers()
-})
+```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

src/cleanup.js

@@ -1,25 +1,22 @@
 import flushMicroTasks from './flush-microtasks'
 
-let internalCleanupCbs = []
-let cleanupCbs = []
+let cleanupCallbacks = []
 
 async function cleanup() {
   await flushMicroTasks()
-  internalCleanupCbs.forEach((cb) => cb())
-  internalCleanupCbs = []
-  cleanupCbs.forEach((cb) => cb())
-}
-
-function addInternalCleanup(callback) {
-  internalCleanupCbs.push(callback)
+  for (const callback of cleanupCallbacks) {
+    await callback()
+  }
+  cleanupCallbacks = []
 }
 
 function addCleanup(callback) {
-  cleanupCbs.push(callback)
+  cleanupCallbacks.unshift(callback)
+  return () => removeCleanup(callback)
 }
 
-function removeInternalCleanup(callback) {
-  internalCleanupCbs = internalCleanupCbs.filter((cb) => cb !== callback)
+function removeCleanup(callback) {
+  cleanupCallbacks = cleanupCallbacks.filter((cb) => cb !== callback)
 }
 
-export { cleanup, addCleanup, addInternalCleanup, removeInternalCleanup }
+export { cleanup, addCleanup, removeCleanup }

src/pure.js

@@ -1,7 +1,7 @@
 import React, { Suspense } from 'react'
 import { act, create } from 'react-test-renderer'
 import asyncUtils from './asyncUtils'
-import { cleanup, addCleanup, addInternalCleanup, removeInternalCleanup } from './cleanup'
+import { cleanup, addCleanup, removeCleanup } from './cleanup'
 
 function TestHook({ callback, hookProps, onError, children }) {
   try {
@@ -84,12 +84,12 @@ function renderHook(callback, { initialProps, wrapper } = {}) {
 
   function unmountHook() {
     act(() => {
-      removeInternalCleanup(unmountHook)
+      removeCleanup(unmountHook)
       unmount()
     })
   }
 
-  addInternalCleanup(unmountHook)
+  addCleanup(unmountHook)
 
   return {
     result,
@@ -99,4 +99,4 @@ function renderHook(callback, { initialProps, wrapper } = {}) {
   }
 }
 
-export { renderHook, cleanup, act, addCleanup }
+export { renderHook, cleanup, addCleanup, removeCleanup, act }

test/addCleanup.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'])
+  })
 })