fix: cleanup async (#1817)

Author: mdjastrzebski

src/__tests__/cleanup.test.tsx

@@ -1,7 +1,7 @@
 import * as React from 'react';
 import { View } from 'react-native';
 
-import { cleanup, render } from '../pure';
+import { cleanupAsync, render, renderAsync } from '../pure';
 
 class Test extends React.Component<{ onUnmount: () => void }> {
   componentWillUnmount() {
@@ -14,13 +14,24 @@ class Test extends React.Component<{ onUnmount: () => void }> {
   }
 }
 
-test('cleanup', () => {
+test('cleanup after render', async () => {
   const fn = jest.fn();
 
   render(<Test onUnmount={fn} />);
   render(<Test onUnmount={fn} />);
   expect(fn).not.toHaveBeenCalled();
 
-  cleanup();
+  await cleanupAsync();
+  expect(fn).toHaveBeenCalledTimes(2);
+});
+
+test('cleanup after renderAsync', async () => {
+  const fn = jest.fn();
+
+  await renderAsync(<Test onUnmount={fn} />);
+  await renderAsync(<Test onUnmount={fn} />);
+  expect(fn).not.toHaveBeenCalled();
+
+  await cleanupAsync();
   expect(fn).toHaveBeenCalledTimes(2);
 });

src/cleanup.ts

@@ -1,8 +1,9 @@
 import { clearRenderResult } from './screen';
 
 type CleanUpFunction = () => void;
+type CleanUpFunctionAsync = () => Promise<void>;
 
-const cleanupQueue = new Set<CleanUpFunction>();
+const cleanupQueue = new Set<CleanUpFunction | CleanUpFunctionAsync>();
 
 export default function cleanup() {
   clearRenderResult();
@@ -11,6 +12,16 @@ export default function cleanup() {
   cleanupQueue.clear();
 }
 
-export function addToCleanupQueue(fn: CleanUpFunction) {
+export async function cleanupAsync() {
+  clearRenderResult();
+
+  for (const fn of cleanupQueue) {
+    await fn();
+  }
+
+  cleanupQueue.clear();
+}
+
+export function addToCleanupQueue(fn: CleanUpFunction | CleanUpFunctionAsync) {
   cleanupQueue.add(fn);
 }

src/index.ts

@@ -3,7 +3,7 @@ import './matchers/extend-expect';
 
 import { getIsReactActEnvironment, setReactActEnvironment } from './act';
 import { flushMicroTasks } from './flush-micro-tasks';
-import { cleanup } from './pure';
+import { cleanupAsync } from './pure';
 
 if (!process?.env?.RNTL_SKIP_AUTO_CLEANUP) {
   // If we're running in a test runner that supports afterEach
@@ -14,7 +14,7 @@ if (!process?.env?.RNTL_SKIP_AUTO_CLEANUP) {
   if (typeof afterEach === 'function') {
     afterEach(async () => {
       await flushMicroTasks();
-      cleanup();
+      await cleanupAsync();
     });
   }
 

src/pure.ts

@@ -1,5 +1,5 @@
 export { default as act } from './act';
-export { default as cleanup } from './cleanup';
+export { default as cleanup, cleanupAsync } from './cleanup';
 export { default as fireEvent, fireEventAsync } from './fire-event';
 export { default as render } from './render';
 export { default as renderAsync } from './render-async';

src/render.tsx

@@ -124,7 +124,7 @@ function buildRenderResult(
     });
   };
 
-  addToCleanupQueue(unmount);
+  addToCleanupQueue(unmountAsync);
 
   const result = {
     ...getQueriesForElement(instance),