From e5ecc520b8e1e6b3bc930984d48cf3ec2fdac3ba Mon Sep 17 00:00:00 2001
From: Robert Knight <robertknight@gmail.com>
Date: Sun, 11 Aug 2019 10:53:37 +0100
Subject: [PATCH] Implement support for nested calls to `act`

When calls to act are nested, effects and component updates are only
flushed when the outer `act` call returns, as per [1] and [2].

This is convenient for creating helper functions which may invoke `act`
themselves.

[1] https://github.com/facebook/react/pull/15682
[2] https://github.com/facebook/react/issues/15472
---
 test-utils/src/index.js            | 37 +++++++++++-
 test-utils/test/shared/act.test.js | 95 +++++++++++++++++++++++++++++-
 2 files changed, 128 insertions(+), 4 deletions(-)

diff --git a/test-utils/src/index.js b/test-utils/src/index.js
index 9e7a677da2c..9d2ee58d15f 100644
--- a/test-utils/src/index.js
+++ b/test-utils/src/index.js
@@ -10,11 +10,37 @@ export function setupRerender() {
 	return () => options.__test__drainQueue && options.__test__drainQueue();
 }
 
+const isThenable = value => value != null && typeof value.then === 'function';
+
+/** Depth of nested calls to `act`. */
+let actDepth = 0;
+
 /**
- * Run a test function, and flush all effects and rerenders after invoking it
- * @param {() => void} cb The function under test
+ * Run a test function, and flush all effects and rerenders after invoking it.
+ *
+ * Returns a Promise which resolves "immediately" if the callback is
+ * synchronous or when the callback's result resolves if it is asynchronous.
+ *
+ * @param {() => void|Promise<void>} cb The function under test. This may be sync or async.
+ * @return {Promise<void>}
  */
 export function act(cb) {
+	++actDepth;
+	if (actDepth > 1) {
+		// If calls to `act` are nested, a flush happens only when the
+		// outermost call returns. In the inner call, we just execute the
+		// callback and return since the infrastructure for flushing has already
+		// been set up.
+		const result = cb();
+		if (isThenable(result)) {
+			return result.then(() => {
+				--actDepth;
+			});
+		}
+		--actDepth;
+		return Promise.resolve();
+	}
+
 	const previousRequestAnimationFrame = options.requestAnimationFrame;
 	const rerender = setupRerender();
 
@@ -36,14 +62,19 @@ export function act(cb) {
 
 		teardown();
 		options.requestAnimationFrame = previousRequestAnimationFrame;
+
+		--actDepth;
 	};
 
 	const result = cb();
 
-	if (result != null && typeof result.then === 'function') {
+	if (isThenable(result)) {
 		return result.then(finish);
 	}
 
+	// nb. If the callback is synchronous, effects must be flushed before
+	// `act` returns, so that the caller does not have to await the result,
+	// even though React recommends this.
 	finish();
 	return Promise.resolve();
 }
diff --git a/test-utils/test/shared/act.test.js b/test-utils/test/shared/act.test.js
index bf9c58a0a6e..296fd64d4bf 100644
--- a/test-utils/test/shared/act.test.js
+++ b/test-utils/test/shared/act.test.js
@@ -1,5 +1,5 @@
 import { options, createElement as h, render } from 'preact';
-import { useEffect, useState } from 'preact/hooks';
+import { useEffect, useReducer, useState } from 'preact/hooks';
 
 import { setupScratch, teardown } from '../../../test/_util/helpers';
 import { act } from '../../src';
@@ -238,4 +238,97 @@ describe('act', () => {
 			'act result resolved'
 		]);
 	});
+
+	context('when `act` calls are nested', () => {
+		it('should invoke nested sync callback and return a Promise', () => {
+			let innerResult;
+			const spy = sinon.stub();
+
+			act(() => {
+				innerResult = act(spy);
+			});
+
+			expect(spy).to.be.calledOnce;
+			expect(innerResult.then).to.be.a('function');
+		});
+
+		it('should invoke nested async callback and return a Promise', async () => {
+			const events = [];
+
+			await act(async () => {
+				events.push('began outer act callback');
+				await act(async () => {
+					events.push('began inner act callback');
+					await Promise.resolve();
+					events.push('end inner act callback');
+				});
+				events.push('end outer act callback');
+			});
+			events.push('act finished');
+
+			expect(events).to.deep.equal([
+				'began outer act callback',
+				'began inner act callback',
+				'end inner act callback',
+				'end outer act callback',
+				'act finished'
+			]);
+		});
+
+		it('should only flush effects when outer `act` call returns', () => {
+			let counter = 0;
+
+			function Widget() {
+				useEffect(() => {
+					++counter;
+				});
+				const [, forceUpdate] = useReducer(x => x + 1, 0);
+				return <button onClick={forceUpdate}>test</button>;
+			}
+
+			act(() => {
+				render(<Widget />, scratch);
+				const button = scratch.querySelector('button');
+				expect(counter).to.equal(0);
+
+				act(() => {
+					button.dispatchEvent(new Event('click'));
+				});
+
+				// Effect triggered by inner `act` call should not have been
+				// flushed yet.
+				expect(counter).to.equal(0);
+			});
+
+			// Effects triggered by inner `act` call should not have been
+			// flushed yet.
+			expect(counter).to.equal(2);
+		});
+
+		it('should only flush updates when outer `act` call returns', () => {
+			function Button() {
+				const [count, setCount] = useState(0);
+				const increment = () => setCount(count => count + 1);
+				return <button onClick={increment}>{count}</button>;
+			}
+
+			render(<Button />, scratch);
+			const button = scratch.querySelector('button');
+			expect(button.textContent).to.equal('0');
+
+			act(() => {
+				act(() => {
+					button.dispatchEvent(new Event('click'));
+				});
+
+				// Update triggered by inner `act` call should not have been
+				// flushed yet.
+				expect(button.textContent).to.equal('0');
+			});
+
+			// Updates from outer and inner `act` calls should now have been
+			// flushed.
+			expect(button.textContent).to.equal('1');
+		});
+	});
 });