test_runner: add cleanup function support to before() hooks

lib/internal/test_runner/test.js

@@ -697,13 +697,15 @@ class Test extends AsyncResource {
     this.subtests = [];
     this.waitingOn = 0;
     this.finished = false;
+    this.pendingBeforeHooks = [];
     this.hooks = {
       __proto__: null,
       before: [],
       after: [],
       beforeEach: [],
       afterEach: [],
       ownAfterEachCount: 0,
+      cleanup: [],
     };
 
     if (loc === undefined) {
@@ -969,12 +971,18 @@ class Test extends AsyncResource {
       hook.run = runOnce(hook.run, kRunOnceOptions);
     }
     if (name === 'before' && this.startTime !== null) {
-      // Test has already started, run the hook immediately
-      PromisePrototypeThen(hook.run(this.getRunArgs()), () => {
+      // Test has already started, run the hook immediately and track the promise
+      // We need to wait for it to complete and capture its return value
+      const hookPromise = (async () => {
+        const result = await hook.run(this.getRunArgs());
+        if (typeof result === 'function') {
+          ArrayPrototypePush(this.hooks.cleanup, result);
+        }
         if (hook.error != null) {
           this.fail(hook.error);
         }
-      });
+      })();
+      ArrayPrototypePush(this.pendingBeforeHooks, hookPromise);
     }
     if (name === 'afterEach') {
       // afterEach hooks for the current test should run in the order that they
@@ -1112,12 +1120,17 @@ class Test extends AsyncResource {
     validateOneOf(hook, 'hook name', kHookNames);
     try {
       const hooks = this.hooks[hook];
+      const cleanup = this.hooks.cleanup;
       for (let i = 0; i < hooks.length; ++i) {
         const hook = hooks[i];
-        await hook.run(args);
+        const hookReturn = await hook.run(args);
         if (hook.error) {
           throw hook.error;
         }
+        if (hook === 'before' && typeof hookReturn === 'function') {
+          ArrayPrototypePush(cleanup, hookReturn);
+        }
+
       }
     } catch (err) {
       const error = new ERR_TEST_FAILURE(`failed running ${hook} hook`, kHookFailure);
@@ -1126,6 +1139,19 @@ class Test extends AsyncResource {
     }
   }
 
+  async runCleanup() {
+    const cleanup = this.hooks.cleanup;
+    // Run cleanup functions in reverse order (LIFO)
+    for (let i = cleanup.length - 1; i >= 0; --i) {
+      const cleanupFn = cleanup[i];
+      try {
+        await cleanupFn();
+      } catch (err) {
+        // Continue with other cleanup functions even if one fails
+      }
+    }
+  }
+
   async filteredRun() {
     this.pass();
     this.subtests = [];
@@ -1166,6 +1192,7 @@ class Test extends AsyncResource {
     }, kRunOnceOptions);
 
     let stopPromise;
+    let testFunctionReturnValue;
 
     try {
       if (this.parent?.hooks.before.length > 0) {
@@ -1205,10 +1232,15 @@ class Test extends AsyncResource {
       ArrayPrototypePush(promises, stopPromise);
 
       // Wait for the race to finish
-      await SafePromiseRace(promises);
+      testFunctionReturnValue = await SafePromiseRace(promises);
 
       this[kShouldAbort]();
 
+      // Wait for any late-registered before hooks to complete
+      if (this.pendingBeforeHooks.length > 0) {
+        await SafePromiseAll(this.pendingBeforeHooks);
+      }
+
       if (this.subtestsPromise !== null) {
         await SafePromiseRace([this.subtestsPromise.promise, stopPromise]);
       }
@@ -1225,6 +1257,9 @@ class Test extends AsyncResource {
       this.pass();
       await afterEach();
       await after();
+      if (this.hooks.cleanup.length > 0) {
+        await this.runCleanup();
+      }
     } catch (err) {
       if (isTestFailureError(err)) {
         if (err.failureType === kTestTimeoutFailure) {
@@ -1282,6 +1317,8 @@ class Test extends AsyncResource {
       await SafePromiseAllReturnVoid(promises);
       process.exit();
     }
+
+    return testFunctionReturnValue;
   }
 
   postRun(pendingSubtestsError) {

test/test-runner/test-before-cleanup.mjs

@@ -0,0 +1,143 @@
+import { describe, test } from 'node:test';
+import assert from 'node:assert/strict';
+
+// Helper to capture execution order
+function makeTracker() {
+  const log = [];
+  return {
+    log,
+    track: (label) => log.push(label),
+    assertOrder: (...expected) => assert.deepEqual(log, expected),
+  };
+}
+
+describe('before hook cleanup function', () => {
+
+  test('cleanup returned from before() is called after the test', async (t) => {
+    const tracker = makeTracker();
+
+    await t.test('run test with cleanup', async (t) => {
+      t.before(() => {
+        tracker.track('before');
+        return () => tracker.track('before-cleanup');
+      });
+
+      t.after(() => tracker.track('after'));
+
+      tracker.track('test-body');
+    });
+
+    // After the nested test completes, verify cleanup ran
+    await t.test('verify cleanup ran', () => {
+      tracker.assertOrder('before', 'test-body', 'after', 'before-cleanup');
+    });
+  });
+
+  test('cleanup returned from before() runs after explicit t.after() hooks', async (t) => {
+    const tracker = makeTracker();
+
+    await t.test('run test with cleanup and after hook', async (t) => {
+      t.before(() => {
+        tracker.track('before');
+        return () => tracker.track('before-cleanup');
+      });
+
+      t.after(() => tracker.track('after'));
+
+      tracker.track('test-body');
+    });
+
+    // Verify cleanup ran after the after hook
+    await t.test('verify order', () => {
+      tracker.assertOrder('before', 'test-body', 'after', 'before-cleanup');
+    });
+  });
+
+  test('multiple before() cleanups run in LIFO order', async (t) => {
+    const tracker = makeTracker();
+
+    await t.test('run test with multiple before hooks', async (t) => {
+      t.before(() => {
+        tracker.track('before-1');
+        return () => tracker.track('cleanup-1');
+      });
+
+      t.before(() => {
+        tracker.track('before-2');
+        return () => tracker.track('cleanup-2');
+      });
+
+      tracker.track('test-body');
+    });
+
+    // Verify cleanups ran in LIFO order
+    await t.test('verify LIFO order', () => {
+      tracker.assertOrder('before-1', 'before-2', 'test-body', 'cleanup-2', 'cleanup-1');
+    });
+  });
+
+  test('before() without a return value does not register a cleanup', async (t) => {
+    const tracker = makeTracker();
+
+    await t.test('run test without cleanup', async (t) => {
+      t.before(() => {
+        tracker.track('before');
+        // no return value
+      });
+
+      t.after(() => tracker.track('after'));
+
+      tracker.track('test-body');
+    });
+
+    // Verify only before, test-body, and after ran (no cleanup)
+    await t.test('verify no cleanup', () => {
+      tracker.assertOrder('before', 'test-body', 'after');
+    });
+  });
+
+  test('async cleanup function from before() works correctly', async (t) => {
+    const tracker = makeTracker();
+
+    await t.test('run test with async cleanup', async (t) => {
+      // Synchronous before hook that returns async cleanup
+      t.before(() => {
+        tracker.track('before');
+        return async () => {
+          await Promise.resolve();
+          tracker.track('async-cleanup');
+        };
+      });
+
+      tracker.track('test-body');
+    });
+
+    // Verify async cleanup ran
+    await t.test('verify async cleanup ran', () => {
+      tracker.assertOrder('before', 'test-body', 'async-cleanup');
+    });
+  });
+
+  test('cleanup from before() runs even if the test body throws', async (t) => {
+    const tracker = makeTracker();
+
+    await t.test('run test that throws', async (t) => {
+      t.before(() => {
+        tracker.track('before');
+        return () => tracker.track('before-cleanup');
+      });
+
+      try {
+        throw new Error('test error');
+      } catch {
+        tracker.track('caught-error');
+      }
+    });
+
+    // Verify cleanup still ran even though test threw
+    await t.test('verify cleanup after error', () => {
+      tracker.assertOrder('before', 'caught-error', 'before-cleanup');
+    });
+  });
+
+});