Add test.cleanup() and test.cleanupEach() for reliable cleanup tasks

docs/01-writing-tests.md

@@ -244,6 +244,38 @@ Keep in mind that the `.beforeEach()` and `.afterEach()` hooks run just before a
 
 Remember that AVA runs each test file in its own process. You may not have to clean up global state in a `.after()`-hook since that's only called right before the process exits.
 
+## Cleanup hooks
+
+It's common to run idempotent cleanup code both before *and* after your tests: before, to remove state left behind by an interrupted or crashed previous run, and after, to remove state created by the current run. `test.cleanup()` and `test.cleanupEach()` let you declare this with a single function.
+
+`test.cleanup()` registers the same implementation as both a `test.before()` hook and a `test.after.always()` hook. It therefore always runs—even when tests fail, or when `--fail-fast` is used. `test.cleanupEach()` does the same, but as a `test.beforeEach()` hook and a `test.afterEach.always()` hook, so the implementation runs around *each* test:
+
+```js
+test.cleanup(() => {
+	if (tempDirExists()) {
+		removeTempDir();
+	}
+});
+
+test.cleanupEach(t => {
+	t.context.db.reset();
+});
+```
+
+Both support the `.skip` modifier, and can be combined with `test.serial`:
+
+```js
+test.serial.cleanup(() => {
+	// Runs as a serial `before` hook and a serial `after.always` hook.
+});
+
+test.cleanup.skip(() => {
+	// Never runs.
+});
+```
+
+
+
 ## Test context
 
 Hooks can share context with the test:

lib/create-chain.js

@@ -152,6 +152,31 @@ export default function createChain(fn, defaults, meta) {
 	root.serial.before = createHookChain(startChain('test.before', fn, {...defaults, serial: true, type: 'before'}), false);
 	root.serial.beforeEach = createHookChain(startChain('test.beforeEach', fn, {...defaults, serial: true, type: 'beforeEach'}), false);
 
+	// `cleanup` registers the same implementation as both a `before` hook and an
+	// `after.always` hook. `cleanupEach` registers it as both a `beforeEach` hook
+	// and an `afterEach.always` hook. The "after" halves always run, even when
+	// tests fail or `--fail-fast` is used, so idempotent cleanup code can rely on
+	// them to remove state left behind by the current (or a previously crashed)
+	// run.
+	function createCleanupChain(node, serial) {
+		const cleanup = startChain(serial ? 'test.serial.cleanup' : 'test.cleanup', fn, {
+			...defaults, serial: Boolean(serial), type: 'cleanup',
+		});
+		extendChain(cleanup, 'skip', 'skipped');
+		node.cleanup = cleanup;
+
+		const cleanupEach = startChain(serial ? 'test.serial.cleanupEach' : 'test.cleanupEach', fn, {
+			...defaults, serial: Boolean(serial), type: 'cleanupEach',
+		});
+		extendChain(cleanupEach, 'skip', 'skipped');
+		node.cleanupEach = cleanupEach;
+
+		return node;
+	}
+
+	createCleanupChain(root, false);
+	createCleanupChain(root.serial, true);
+
 	// "todo" tests cannot be chained. Allow todo tests to be flagged as needing
 	// to be serial.
 	root.todo = startChain('test.todo', fn, {...defaults, type: 'test', todo: true});

lib/runner.js

@@ -174,7 +174,18 @@ export default class Runner extends Emittery {
 						todo: false,
 					});
 				} else if (!metadata.skipped) {
-					this.tasks[metadata.type + (metadata.always ? 'Always' : '')].push(task);
+					if (metadata.type === 'cleanup' || metadata.type === 'cleanupEach') {
+						// `cleanup` runs as both a `before` hook and an `after.always` hook.
+						// `cleanupEach` runs as both a `beforeEach` hook and an
+						// `afterEach.always` hook. The "after" halves always run, even when
+						// tests fail or `--fail-fast` is used.
+						const beforeType = metadata.type === 'cleanup' ? 'before' : 'beforeEach';
+						const afterType = metadata.type === 'cleanup' ? 'afterAlways' : 'afterEachAlways';
+						this.tasks[beforeType].push(task);
+						this.tasks[afterType].push({...task, metadata: {...task.metadata, always: true}});
+					} else {
+						this.tasks[metadata.type + (metadata.always ? 'Always' : '')].push(task);
+					}
 				}
 			}
 		}, {

test-types/module/conditional-chains.ts

@@ -10,5 +10,14 @@ anyTest.skipIf(true).todo('skipIf todo should be allowed');
 anyTest.runIf(false).todo('runIf todo should be allowed');
 anyTest.skipIf(true).serial.todo('skipIf serial todo should be allowed');
 anyTest.runIf(false).serial.todo('runIf serial todo should be allowed');
+
+anyTest.cleanup(() => {});
+anyTest.cleanupEach(() => {});
+anyTest.cleanup.skip(() => {});
+anyTest.cleanupEach.skip(() => {});
+anyTest.serial.cleanup(() => {});
+anyTest.serial.cleanupEach(() => {});
+anyTest.skipIf(true).cleanup(() => {});
+anyTest.runIf(false).cleanupEach(() => {});
 anyTest.serial.skipIf(true).todo('serial skipIf todo should be allowed');
 anyTest.serial.runIf(false).todo('serial runIf todo should be allowed');

test/cleanup/test.js

@@ -0,0 +1,62 @@
+import test from 'ava';
+import fs from 'node:fs';
+import os from 'node:os';
+import path from 'node:path';
+
+const marker = path.join(os.tmpdir(), `ava-cleanup-${process.pid}.log`);
+try {
+	fs.unlinkSync(marker);
+} catch {}
+
+// `test.cleanup()` runs as both a `before` hook and an `after.always` hook.
+test.cleanup(() => {
+	fs.appendFileSync(marker, 'x');
+});
+
+let testCount = 0;
+
+test('the before-half of cleanup runs before tests', t => {
+	testCount++;
+	t.true(fs.existsSync(marker));
+});
+
+test.after.always(() => {
+	const content = fs.readFileSync(marker, 'utf8');
+	// `cleanup` runs once as a `before` hook and once as an `after.always` hook,
+	// regardless of how many tests there are.
+	if (content !== 'xx') {
+		throw new Error('expected cleanup to run as both a before and after.always hook, but got ' + JSON.stringify(content));
+	}
+});
+
+const markerEach = path.join(os.tmpdir(), `ava-cleanup-each-${process.pid}.log`);
+try {
+	fs.unlinkSync(markerEach);
+} catch {}
+
+// `test.cleanupEach()` runs as both a `beforeEach` hook and an
+// `afterEach.always` hook, so it runs around each test.
+test.cleanupEach(() => {
+	fs.appendFileSync(markerEach, 'y');
+});
+
+test('cleanupEach runs around the first test', t => {
+	testCount++;
+	t.true(fs.existsSync(markerEach));
+});
+
+test('cleanupEach runs around the second test', t => {
+	testCount++;
+	t.true(fs.existsSync(markerEach));
+});
+
+test.after.always(() => {
+	const content = fs.readFileSync(markerEach, 'utf8');
+	// `cleanupEach` wraps every test as a `beforeEach` + `afterEach.always` hook,
+	// so it runs twice per test. `testCount` counts the tests above (the cleanup
+	// test and these two), giving the total number of wrapped tests.
+	const expected = testCount * 2;
+	if (content.length !== expected || ![...content].every(c => c === 'y')) {
+		throw new Error('expected ' + expected + ' cleanupEach executions, but got ' + JSON.stringify(content));
+	}
+});

test/create-chain/test.js

@@ -245,3 +245,66 @@ test('skipIf(false).runIf(true) does not skip', t => {
 	t.is(calls.length, 1);
 	t.is(calls[0].metadata.skipped, undefined);
 });
+
+test('cleanup() registers a cleanup task', t => {
+	const {calls, chain} = createTestChain();
+
+	t.notThrows(() => {
+		chain.cleanup('title', () => {});
+		chain.cleanup(() => {});
+	});
+
+	t.is(calls.length, 2);
+	t.is(calls[0].metadata.type, 'cleanup');
+	t.is(calls[1].metadata.type, 'cleanup');
+});
+
+test('cleanup.skip() skips the cleanup task', t => {
+	const {calls, chain} = createTestChain();
+
+	chain.cleanup.skip('title', () => {});
+
+	t.is(calls.length, 1);
+	t.is(calls[0].metadata.type, 'cleanup');
+	t.is(calls[0].metadata.skipped, true);
+});
+
+test('cleanupEach() registers a cleanupEach task', t => {
+	const {calls, chain} = createTestChain();
+
+	t.notThrows(() => {
+		chain.cleanupEach('title', () => {});
+		chain.cleanupEach(() => {});
+	});
+
+	t.is(calls.length, 2);
+	t.is(calls[0].metadata.type, 'cleanupEach');
+	t.is(calls[1].metadata.type, 'cleanupEach');
+});
+
+test('serial.cleanup() preserves the serial flag', t => {
+	const {calls, chain} = createTestChain();
+
+	chain.serial.cleanup('title', () => {});
+
+	t.is(calls.length, 1);
+	t.is(calls[0].metadata.type, 'cleanup');
+	t.is(calls[0].metadata.serial, true);
+});
+
+test('serial.cleanupEach() preserves the serial flag', t => {
+	const {calls, chain} = createTestChain();
+
+	chain.serial.cleanupEach('title', () => {});
+
+	t.is(calls.length, 1);
+	t.is(calls[0].metadata.type, 'cleanupEach');
+	t.is(calls[0].metadata.serial, true);
+});
+
+test('cleanup() and cleanupEach() are reachable from conditional chains', t => {
+	const {chain} = createTestChain();
+
+	t.is(typeof chain.skipIf(false).cleanup, 'function');
+	t.is(typeof chain.runIf(true).cleanupEach, 'function');
+});

types/test-fn.d.ts

@@ -87,6 +87,8 @@ export type TestFn<Context = unknown> = {
 	afterEach: AfterFn<Context>;
 	before: BeforeFn<Context>;
 	beforeEach: BeforeFn<Context>;
+	cleanup: CleanupFn<Context>;
+	cleanupEach: CleanupFn<Context>;
 	failing: FailingFn<Context>;
 	macro: MacroFn<Context>;
 	meta: Meta;
@@ -150,6 +152,23 @@ export type BeforeFn<Context = unknown> = {
 	skip: HookSkipFn<Context>;
 };
 
+export type CleanupFn<Context = unknown> = {
+	/**
+	 * Declare a cleanup task. The implementation runs as both a `before` hook and
+	 * an `after.always` hook, so it always runs—even when tests fail or
+	 * `--fail-fast` is used. Additional arguments are passed to the implementation or macro.
+	 */
+	<Args extends unknown[]>(title: string, implementation: Implementation<Args, Context>, ...args: Args): void;
+
+	/**
+	 * Declare a cleanup task. The implementation runs as both a `before` hook and
+	 * an `after.always` hook. Additional arguments are passed to the implementation or macro.
+	 */
+	<Args extends unknown[]>(implementation: Implementation<Args, Context>, ...args: Args): void;
+
+	skip: HookSkipFn<Context>;
+};
+
 export type FailingFn<Context = unknown> = {
 	/**
 	 * Declare a concurrent test that is expected to fail.
@@ -209,6 +228,8 @@ export type SerialFn<Context = unknown> = {
 	afterEach: AfterFn<Context>;
 	before: BeforeFn<Context>;
 	beforeEach: BeforeFn<Context>;
+	cleanup: CleanupFn<Context>;
+	cleanupEach: CleanupFn<Context>;
 	failing: FailingFn<Context>;
 	only: OnlyFn<Context>;
 	/** Declare a test that only runs when `condition` is true; otherwise the test is skipped. */