Implement registerCompletionHandler()

Register a function to be called when AVA has completed a test run without uncaught exceptions or unhandled rejections.

Fixes #3279.
 *
 * Completion handlers are invoked in order of registration. Results are not awaited.

Author: novemberborn

docs/01-writing-tests.md

@@ -154,7 +154,7 @@ AVA lets you register hooks that are run before and after your tests. This allow
 
 If a test is skipped with the `.skip` modifier, the respective `.beforeEach()`, `.afterEach()` and `.afterEach.always()` hooks are not run. Likewise, if all tests in a test file are skipped `.before()`, `.after()` and `.after.always()` hooks for the file are not run.
 
-*You may not need to use `.afterEach.always()` hooks to clean up after a test.* You can use [`t.teardown()`](./02-execution-context.md#tteardownfn) to undo side-effects *within* a particular test.
+*You may not need to use `.afterEach.always()` hooks to clean up after a test.* You can use [`t.teardown()`](./02-execution-context.md#tteardownfn) to undo side-effects *within* a particular test. Or use [`registerCompletionHandler()`](./08-common-pitfalls.md#timeouts-because-a-file-failed-to-exit) to run cleanup code after AVA has completed its work.
 
 Like `test()` these methods take an optional title and an implementation function. The title is shown if your hook fails to execute. The implementation is called with an [execution object](./02-execution-context.md). You can use assertions in your hooks. You can also pass a [macro function](#reusing-test-logic-through-macros) and additional arguments.
 

docs/07-test-timeouts.md

@@ -4,7 +4,7 @@ Translations: [Français](https://github.com/avajs/ava-docs/blob/main/fr_FR/docs
 
 [![Open in StackBlitz](https://developer.stackblitz.com/img/open_in_stackblitz.svg)](https://stackblitz.com/github/avajs/ava/tree/main/examples/timeouts?file=test.js&terminal=test&view=editor)
 
-Timeouts in AVA behave differently than in other test frameworks. AVA resets a timer after each test, forcing tests to quit if no new test results were received within the specified timeout. This can be used to handle stalled tests.
+Timeouts in AVA behave differently than in other test frameworks. AVA resets a timer after each test, forcing tests to quit if no new test results were received within the specified timeout. This can be used to handle stalled tests. This same mechanism is used to determine when a test file is preventing a clean exit.
 
 The default timeout is 10 seconds.
 

docs/08-common-pitfalls.md

@@ -81,6 +81,37 @@ Error [ERR_WORKER_INVALID_EXEC_ARGV]: Initiated Worker with invalid execArgv fla
 
 If possible don't specify the command line option when running AVA. Alternatively you could [disable worker threads in AVA](./06-configuration.md#options).
 
+## Timeouts because a file failed to exit
+
+You may get a "Timed out while running tests" error because AVA failed to exit when running a particular file.
+
+AVA waits for Node.js to exit the worker thread or child process. If this takes too long, AVA counts it as a timeout.
+
+It is best practice to make sure your code exits cleanly. We've also seen occurrences where an explicit `process.exit()` call inside a worker thread could not be observed in AVA's main process.
+
+For these reasons we're not providing an option to disable this timeout behavior. However, it is possible to register a callback for when AVA has completed the test run without uncaught exceptions or unhandled rejections. From inside this callback you can do whatever you need to do, including calling `process.exit()`.
+
+Create a `_force-exit.mjs` file:
+
+```js
+import process from 'node:process';
+import { registerCompletionHandler } from 'ava';
+
+registerCompletionHandler(() => {
+	process.exit();
+});
+```
+
+Completion handlers are invoked in order of registration. Results are not awaited.
+
+Load it for all test files through AVA's `require` option:
+
+```js
+export default {
+	require: ['./_force-exit.mjs'],
+};
+```
+
 ## Sharing variables between asynchronous tests
 
 By default AVA executes tests concurrently. This can cause problems if your tests are asynchronous and share variables.

entrypoints/main.d.mts

@@ -10,3 +10,11 @@ declare const test: TestFn;
 
 /** Call to declare a test, or chain to declare hooks or test modifiers */
 export default test;
+
+/**
+ * Register a function to be called when AVA has completed a test run without uncaught exceptions or unhandled rejections.
+ *
+ * Completion handlers are invoked in order of registration. Results are not awaited.
+ */
+declare const registerCompletionHandler: (handler: () => void) => void;
+export {registerCompletionHandler};

entrypoints/main.mjs

@@ -1 +1,2 @@
 export {default} from '../lib/worker/main.cjs';
+export {registerCompletionHandler} from '../lib/worker/completion-handlers.js';

lib/worker/base.js

@@ -15,6 +15,7 @@ import Runner from '../runner.js';
 import serializeError from '../serialize-error.js';
 
 import channel from './channel.cjs';
+import {runCompletionHandlers} from './completion-handlers.js';
 import lineNumberSelection from './line-numbers.js';
 import {set as setOptions} from './options.cjs';
 import {flags, refs, sharedWorkerTeardowns} from './state.cjs';
@@ -23,17 +24,22 @@ import {isRunningInThread, isRunningInChildProcess} from './utils.cjs';
 const currentlyUnhandled = setUpCurrentlyUnhandled();
 let runner;
 
-let forcingExit = false;
+let expectingExit = false;
 
 const forceExit = () => {
-	forcingExit = true;
+	expectingExit = true;
 	process.exit(1);
 };
 
+const avaIsDone = () => {
+	expectingExit = true;
+	runCompletionHandlers();
+};
+
 // Override process.exit with an undetectable replacement
 // to report when it is called from a test (which it should never be).
 const handleProcessExit = (target, thisArg, args) => {
-	if (!forcingExit) {
+	if (!expectingExit) {
 		const error = new Error('Unexpected process.exit()');
 		Error.captureStackTrace(error, handleProcessExit);
 		channel.send({type: 'process-exit', stack: error.stack});
@@ -118,7 +124,7 @@ const run = async options => {
 		nowAndTimers.setImmediate(() => {
 			const unhandled = currentlyUnhandled();
 			if (unhandled.length === 0) {
-				return;
+				return avaIsDone();
 			}
 
 			for (const rejection of unhandled) {

lib/worker/completion-handlers.js

@@ -0,0 +1,13 @@
+import process from 'node:process';
+
+import state from './state.cjs';
+
+export function runCompletionHandlers() {
+	for (const handler of state.completionHandlers) {
+		process.nextTick(() => handler());
+	}
+}
+
+export function registerCompletionHandler(handler) {
+	state.completionHandlers.push(handler);
+}

lib/worker/state.cjs

@@ -1,5 +1,6 @@
 'use strict';
 exports.flags = {loadedMain: false};
 exports.refs = {runnerChain: null};
+exports.completionHandlers = [];
 exports.sharedWorkerTeardowns = [];
 exports.waitForReady = [];

test/completion-handlers/fixtures/exit0.js

@@ -0,0 +1,7 @@
+import test, { registerCompletionHandler } from 'ava'
+
+registerCompletionHandler(() => {
+	process.exit(0)
+})
+
+test('pass', t => t.pass())

test/completion-handlers/fixtures/one.js

@@ -0,0 +1,9 @@
+import test, { registerCompletionHandler } from 'ava'
+
+registerCompletionHandler(() => {
+	console.error('one')
+})
+
+test('pass', t => {
+	t.pass()
+})

test/completion-handlers/fixtures/package.json

@@ -0,0 +1,8 @@
+{
+	"type": "module",
+	"ava": {
+		"files": [
+			"*.js"
+		]
+	}
+}

test/completion-handlers/fixtures/two.js

@@ -0,0 +1,10 @@
+import test, { registerCompletionHandler } from 'ava'
+
+registerCompletionHandler(() => {
+	console.error('one')
+})
+registerCompletionHandler(() => {
+	console.error('two')
+})
+
+test('pass', t => t.pass())

test/completion-handlers/test.js

@@ -0,0 +1,17 @@
+import test from '@ava/test';
+
+import {cleanOutput, fixture} from '../helpers/exec.js';
+
+test('runs a single completion handler', async t => {
+	const result = await fixture(['one.js']);
+	t.is(cleanOutput(result.stderr), 'one');
+});
+
+test('runs multiple completion handlers in registration order', async t => {
+	const result = await fixture(['two.js']);
+	t.deepEqual(cleanOutput(result.stderr).split('\n'), ['one', 'two']);
+});
+
+test('completion handlers may exit the process', async t => {
+	await t.notThrowsAsync(fixture(['exit0.js']));
+});