Implement support for async act callbacks and nested calls to act

.babelrc

@@ -18,6 +18,7 @@
   ],
   "plugins": [
     "transform-object-rest-spread",
-    "transform-react-jsx"
+    "transform-react-jsx",
+    "transform-async-to-promises"
   ]
 }

package.json

@@ -113,6 +113,7 @@
     "babel-core": "6.26.3",
     "babel-loader": "7.1.5",
     "babel-plugin-istanbul": "5.0.1",
+    "babel-plugin-transform-async-to-promises": "^0.8.14",
     "babel-plugin-transform-object-rest-spread": "^6.23.0",
     "babel-plugin-transform-react-jsx": "^6.24.1",
     "babel-preset-env": "^1.6.1",

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();
 
@@ -24,20 +50,33 @@ export function act(cb) {
 	// Override requestAnimationFrame so we can flush pending hooks.
 	options.requestAnimationFrame = (fc) => flush = fc;
 
-	// Execute the callback we were passed.
-	cb();
-	rerender();
+	const finish = () => {
+		rerender();
+		while (flush) {
+			toFlush = flush;
+			flush = null;
 
-	while (flush) {
-		toFlush = flush;
-		flush = null;
+			toFlush();
+			rerender();
+		}
 
-		toFlush();
-		rerender();
+		teardown();
+		options.requestAnimationFrame = previousRequestAnimationFrame;
+
+		--actDepth;
+	};
+
+	const result = cb();
+
+	if (isThenable(result)) {
+		return result.then(finish);
 	}
 
-	teardown();
-	options.requestAnimationFrame = previousRequestAnimationFrame;
+	// 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();
 }
 
 /**

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';
@@ -191,4 +191,144 @@ describe('act', () => {
 		});
 		expect(scratch.firstChild.textContent).to.equal('1');
 	});
+
+	it('returns a Promise if invoked with a sync callback', () => {
+		const result = act(() => {});
+		expect(result.then).to.be.a('function');
+		return result;
+	});
+
+	it('returns a Promise if invoked with an async callback', () => {
+		const result = act(async () => {});
+		expect(result.then).to.be.a('function');
+		return result;
+	});
+
+	it('should await "thenable" result of callback before flushing', async () => {
+		const events = [];
+
+		function TestComponent() {
+			useEffect(() => {
+				events.push('flushed effect');
+			}, []);
+			events.push('scheduled effect');
+			return <div>Test</div>;
+		}
+
+		const delay = ms => new Promise(resolve => setTimeout(resolve, ms));
+
+		events.push('began test');
+		const acted = act(async () => {
+			events.push('began act callback');
+			await delay(1);
+			render(<TestComponent />, scratch);
+			events.push('end act callback');
+		});
+		events.push('act returned');
+		await acted;
+		events.push('act result resolved');
+
+		expect(events).to.deep.equal([
+			'began test',
+			'began act callback',
+			'act returned',
+			'scheduled effect',
+			'end act callback',
+			'flushed effect',
+			'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 now have been
+			// flushed.
+			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');
+		});
+	});
 });