[jest] fix and improve hooks (#1689)

- wait for async hooks to complete before running tests
- add support for `done(err)` callbacks in hooks

fixes #1688

4 files changed, 189 insertions, 8 deletions

diff --git a/packages/bun-types/bun-test.d.ts b/packages/bun-types/bun-test.d.ts
index 0d643ea56..7567c5ad2 100644
--- a/packages/bun-types/bun-test.d.ts
+++ b/packages/bun-types/bun-test.d.ts
@@ -26,11 +26,11 @@ declare module "bun:test" {
   export { test as it };
 
   export function expect(value: any): Expect;
-  export function afterAll(fn: () => void): void;
-  export function beforeAll(fn: () => void): void;
+  export function afterAll(fn: (done: (err?: any) => void) => void | Promise<any>): void;
+  export function beforeAll(fn: (done: (err?: any) => void) => void | Promise<any>): void;
 
-  export function afterEach(fn: () => void): void;
-  export function beforeEach(fn: () => void): void;
+  export function afterEach(fn: (done: (err?: any) => void) => void | Promise<any>): void;
+  export function beforeEach(fn: (done: (err?: any) => void) => void | Promise<any>): void;
 
   interface Expect {
     not: Expect;
diff --git a/packages/bun-types/globals.d.ts b/packages/bun-types/globals.d.ts
index aa532bd42..a27ee2b92 100644
--- a/packages/bun-types/globals.d.ts
+++ b/packages/bun-types/globals.d.ts
@@ -1281,6 +1281,14 @@ declare function queueMicrotask(callback: (...args: any[]) => void): void;
  */
 declare function reportError(error: any): void;
 /**
+ * Run a function immediately after main event loop is vacant
+ * @param handler function to call
+ */
+declare function setImmediate(
+  handler: TimerHandler,
+  ...arguments: any[]
+): number;
+/**
  * Run a function every `interval` milliseconds
  * @param handler function to call
  * @param interval milliseconds to wait between calls
@@ -2277,7 +2285,7 @@ declare function alert(message?: string): void;
 declare function confirm(message?: string): boolean;
 declare function prompt(message?: string, _default?: string): string | null;
 
-/* 
+/*
 
  Web Crypto API
 
diff --git a/src/bun.js/test/jest.zig b/src/bun.js/test/jest.zig
index 2a98d42f8..230c3a96e 100644
--- a/src/bun.js/test/jest.zig
+++ b/src/bun.js/test/jest.zig
@@ -1398,6 +1398,7 @@ pub const DescribeScope = struct {
     file_id: TestRunner.File.ID,
     current_test_id: TestRunner.Test.ID = 0,
     value: JSValue = .zero,
+    done: bool = false,
 
     pub fn push(new: *DescribeScope) void {
         if (comptime is_bindgen) return undefined;
@@ -1486,16 +1487,61 @@ pub const DescribeScope = struct {
         },
     );
 
+    pub fn onDone(
+        ctx: js.JSContextRef,
+        callframe: *JSC.CallFrame,
+    ) callconv(.C) JSValue {
+        const function = callframe.callee();
+        const args = callframe.arguments(1);
+        defer ctx.bunVM().autoGarbageCollect();
+
+        if (JSC.getFunctionData(function)) |data| {
+            var scope = bun.cast(*DescribeScope, data);
+            JSC.setFunctionData(function, null);
+            if (args.len > 0) {
+                const err = args.ptr[0];
+                ctx.bunVM().runErrorHandlerWithDedupe(err, null);
+            }
+            scope.done = true;
+        }
+
+        return JSValue.jsUndefined();
+    }
+
     pub fn execCallback(this: *DescribeScope, ctx: js.JSContextRef, comptime hook: LifecycleHook) JSValue {
         const name = comptime @as(string, @tagName(hook));
         var hooks: []JSC.JSValue = @field(this, name).items;
         for (hooks) |cb, i| {
             if (cb.isEmpty()) continue;
 
-            const err = cb.call(ctx, &.{});
-            if (err.isAnyError(ctx)) {
-                return err;
+            const pending_test = Jest.runner.?.pending_test;
+            // forbid `expect()` within hooks
+            Jest.runner.?.pending_test = null;
+            var result = if (cb.getLengthOfArray(ctx) > 0) brk: {
+                this.done = false;
+                const done_func = JSC.NewFunctionWithData(
+                    ctx,
+                    ZigString.static("done"),
+                    0,
+                    DescribeScope.onDone,
+                    false,
+                    this,
+                );
+                var result = cb.call(ctx, &.{done_func});
+                var vm = VirtualMachine.get();
+                while (!this.done) {
+                    vm.eventLoop().autoTick();
+                    if (this.done) break;
+                    vm.eventLoop().tick();
+                }
+                break :brk result;
+            } else cb.call(ctx, &.{});
+            if (result.asPromise()) |promise| {
+                VirtualMachine.get().waitForPromise(promise);
+                result = promise.result(ctx.vm());
             }
+            Jest.runner.?.pending_test = pending_test;
+            if (result.isAnyError(ctx)) return result;
 
             if (comptime hook == .beforeAll or hook == .afterAll) {
                 hooks[i] = JSC.JSValue.zero;
diff --git a/test/bun.js/bun-test/jest-hooks.test.ts b/test/bun.js/bun-test/jest-hooks.test.ts
index a75025041..4675acb07 100644
--- a/test/bun.js/bun-test/jest-hooks.test.ts
+++ b/test/bun.js/bun-test/jest-hooks.test.ts
@@ -73,4 +73,131 @@ describe("test jest hooks in bun-test", () => {
       expect(animal).toEqual("lion");
     });
   });
+
+  describe("test async hooks", async () => {
+    let beforeAllCalled = 0;
+    let beforeEachCalled = 0;
+    let afterAllCalled = 0;
+    let afterEachCalled = 0;
+
+    beforeAll(async () => {
+      beforeAllCalled += await 1;
+    });
+
+    beforeEach(async () => {
+      beforeEachCalled += await 1;
+    });
+
+    afterAll(async () => {
+      afterAllCalled += await 1;
+    });
+
+    afterEach(async () => {
+      afterEachCalled += await 1;
+    });
+
+    it("should run after beforeAll()", () => {
+      expect(beforeAllCalled).toBe(1);
+      expect(beforeEachCalled).toBe(1);
+      expect(afterAllCalled).toBe(0);
+      expect(afterEachCalled).toBe(0);
+    });
+
+    it("should run after beforeEach()", () => {
+      expect(beforeAllCalled).toBe(1);
+      expect(beforeEachCalled).toBe(2);
+      expect(afterAllCalled).toBe(0);
+      expect(afterEachCalled).toBe(1);
+    });
+  });
+
+  describe("test done callback in hooks", () => {
+    let beforeAllCalled = 0;
+    let beforeEachCalled = 0;
+    let afterAllCalled = 0;
+    let afterEachCalled = 0;
+
+    beforeAll(done => {
+      setImmediate(() => {
+        beforeAllCalled++;
+        done();
+      });
+    });
+
+    beforeEach(done => {
+      setImmediate(() => {
+        beforeEachCalled++;
+        done();
+      });
+    });
+
+    afterAll(done => {
+      setImmediate(() => {
+        afterAllCalled++;
+        done();
+      });
+    });
+
+    afterEach(done => {
+      setImmediate(() => {
+        afterEachCalled++;
+        done();
+      });
+    });
+
+    it("should run after beforeAll()", () => {
+      expect(beforeAllCalled).toBe(1);
+      expect(beforeEachCalled).toBe(1);
+      expect(afterAllCalled).toBe(0);
+      expect(afterEachCalled).toBe(0);
+    });
+
+    it("should run after beforeEach()", () => {
+      expect(beforeAllCalled).toBe(1);
+      expect(beforeEachCalled).toBe(2);
+      expect(afterAllCalled).toBe(0);
+      expect(afterEachCalled).toBe(1);
+    });
+  });
+
+  describe("test async hooks with done()", () => {
+    let beforeAllCalled = 0;
+    let beforeEachCalled = 0;
+    let afterAllCalled = 0;
+    let afterEachCalled = 0;
+
+    beforeAll(async done => {
+      beforeAllCalled += await 1;
+      setTimeout(done, 1);
+    });
+
+    beforeEach(async done => {
+      beforeEachCalled += await 1;
+      setTimeout(done, 1);
+    });
+
+    afterAll(async done => {
+      afterAllCalled += await 1;
+      setTimeout(done, 1);
+    });
+
+    afterEach(async done => {
+      afterEachCalled += await 1;
+      setTimeout(done, 1);
+    });
+
+    it("should run after beforeAll()", () => {
+      expect(beforeAllCalled).toBe(1);
+      expect(beforeEachCalled).toBe(1);
+      expect(afterAllCalled).toBe(0);
+      expect(afterEachCalled).toBe(0);
+    });
+
+    it("should run after beforeEach()", () => {
+      expect(beforeAllCalled).toBe(1);
+      expect(beforeEachCalled).toBe(2);
+      expect(afterAllCalled).toBe(0);
+      expect(afterEachCalled).toBe(1);
+    });
+  });
 });