From aac2e44fc8b970deed16248913582459cf044c64 Mon Sep 17 00:00:00 2001
From: Kamran Ayub <kamran.ayub@gmail.com>
Date: Wed, 24 Jun 2020 23:21:29 -0500
Subject: [PATCH] fix(cleanup): Cleanup should flush microtask queue after
 unmount (#632)

---
 src/__tests__/cleanup.js | 14 ++++++++++++++
 src/pure.js              |  4 +++-
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/src/__tests__/cleanup.js b/src/__tests__/cleanup.js
index d1332bf4..f3f6a557 100644
--- a/src/__tests__/cleanup.js
+++ b/src/__tests__/cleanup.js
@@ -26,3 +26,17 @@ test('cleanup does not error when an element is not a child', async () => {
   render(<div />, {container: document.createElement('div')})
   await cleanup()
 })
+
+test('cleanup waits for queued microtasks during unmount sequence', async () => {
+  const spy = jest.fn()
+
+  const Test = () => {
+    React.useEffect(() => () => setImmediate(spy))
+
+    return null
+  }
+
+  render(<Test />)
+  await cleanup()
+  expect(spy).toHaveBeenCalledTimes(1)
+})
diff --git a/src/pure.js b/src/pure.js
index fe7e6cf9..be0fa3f1 100644
--- a/src/pure.js
+++ b/src/pure.js
@@ -97,8 +97,10 @@ function render(
 }
 
 async function cleanup() {
-  await flush()
   mountedContainers.forEach(cleanupAtContainer)
+  // flush microtask queue after unmounting in case
+  // unmount sequence generates new microtasks
+  await flush()
 }
 
 // maybe one day we'll expose this (perhaps even as a utility returned by render).