feat(bun:test) add support for test.each() and describe.each() (#4047)

* rename callback to func

* update testscope to handle function arguments

* works

* big cleanup

* works in debug, not release

* fix memory issue & update tests

* catch & str test

* write types for each() & switch tests to ts

* rm & typo

* move some code around & support describe

* review changes

3 files changed, 324 insertions, 16 deletions

diff --git a/packages/bun-types/bun-test.d.ts b/packages/bun-types/bun-test.d.ts
index f6e968906..790d8bfea 100644
--- a/packages/bun-types/bun-test.d.ts
+++ b/packages/bun-types/bun-test.d.ts
@@ -169,6 +169,25 @@ declare module "bun:test" {
      * @param condition if these tests should be skipped
      */
     skipIf(condition: boolean): (label: string, fn: () => void) => void;
+    /**
+     * Returns a function that runs for each item in `table`.
+     *
+     * @param table Array of Arrays with the arguments that are passed into the test fn for each row.
+     */
+    each<T extends ReadonlyArray<unknown>>(
+      table: ReadonlyArray<T>,
+    ): (
+      label: string,
+      fn: (...args: T) => void | Promise<unknown>,
+      options?: number | TestOptions,
+    ) => void;
+    each<T>(
+      table: ReadonlyArray<T>,
+    ): (
+      label: string,
+      fn: (arg: T) => void | Promise<unknown>,
+      options?: number | TestOptions,
+    ) => void;
   };
   /**
    * Describes a group of related tests.
@@ -395,6 +414,25 @@ declare module "bun:test" {
         | ((done: (err?: unknown) => void) => void),
       options?: number | TestOptions,
     ) => void;
+    /**
+     * Returns a function that runs for each item in `table`.
+     *
+     * @param table Array of Arrays with the arguments that are passed into the test fn for each row.
+     */
+    each<T extends ReadonlyArray<unknown>>(
+      table: ReadonlyArray<T>,
+    ): (
+      label: string,
+      fn: (...args: T) => void | Promise<unknown>,
+      options?: number | TestOptions,
+    ) => void;
+    each<T>(
+      table: ReadonlyArray<T>,
+    ): (
+      label: string,
+      fn: (arg: T, done: (err?: unknown) => void) => void | Promise<unknown>,
+      options?: number | TestOptions,
+    ) => void;
   };
   /**
    * Runs a test.
diff --git a/src/bun.js/test/jest.zig b/src/bun.js/test/jest.zig
index 8691e5a2d..9169c994b 100644
--- a/src/bun.js/test/jest.zig
+++ b/src/bun.js/test/jest.zig
@@ -373,6 +373,11 @@ pub const Jest = struct {
             ZigString.static("skipIf"),
             JSC.NewFunction(globalObject, ZigString.static("skipIf"), 2, TestScope.skipIf, false),
         );
+        test_fn.put(
+            globalObject,
+            ZigString.static("each"),
+            JSC.NewFunction(globalObject, ZigString.static("each"), 2, TestScope.each, false),
+        );
 
         module.put(
             globalObject,
@@ -405,6 +410,11 @@ pub const Jest = struct {
             ZigString.static("skipIf"),
             JSC.NewFunction(globalObject, ZigString.static("skipIf"), 2, DescribeScope.skipIf, false),
         );
+        describe.put(
+            globalObject,
+            ZigString.static("each"),
+            JSC.NewFunction(globalObject, ZigString.static("each"), 2, DescribeScope.each, false),
+        );
 
         module.put(
             globalObject,
@@ -546,7 +556,11 @@ pub const Jest = struct {
 pub const TestScope = struct {
     label: string = "",
     parent: *DescribeScope,
-    callback: JSC.JSValue,
+
+    func: JSC.JSValue,
+    func_arg: []JSC.JSValue,
+    func_has_callback: bool = false,
+
     id: TestRunner.Test.ID = 0,
     promise: ?*JSInternalPromise = null,
     ran: bool = false,
@@ -578,6 +592,10 @@ pub const TestScope = struct {
         return createScope(globalThis, callframe, "test.todo()", true, .todo);
     }
 
+    pub fn each(globalThis: *JSGlobalObject, callframe: *CallFrame) callconv(.C) JSValue {
+        return createEach(globalThis, callframe, "test.each()", "each", true);
+    }
+
     pub fn callIf(globalThis: *JSGlobalObject, callframe: *CallFrame) callconv(.C) JSValue {
         return createIfScope(globalThis, callframe, "test.if()", "if", TestScope, false);
     }
@@ -637,17 +655,19 @@ pub const TestScope = struct {
     ) Result {
         if (comptime is_bindgen) return undefined;
         var vm = VirtualMachine.get();
-        const callback = this.callback;
+        const func = this.func;
         Jest.runner.?.did_pending_test_fail = false;
         defer {
-            callback.unprotect();
-            this.callback = .zero;
+            for (this.func_arg) |arg| {
+                arg.unprotect();
+            }
+            func.unprotect();
+            this.func = .zero;
+            this.func_has_callback = false;
             vm.autoGarbageCollect();
         }
         JSC.markBinding(@src());
 
-        const callback_length = callback.getLength(vm.global);
-
         var initial_value = JSValue.zero;
         if (test_elapsed_timer) |timer| {
             timer.reset();
@@ -659,7 +679,7 @@ pub const TestScope = struct {
             task.test_id,
         );
 
-        if (callback_length > 0) {
+        if (this.func_has_callback) {
             const callback_func = JSC.NewFunctionWithData(
                 vm.global,
                 ZigString.static("done"),
@@ -669,11 +689,11 @@ pub const TestScope = struct {
                 task,
             );
             task.done_callback_state = .pending;
-            initial_value = callback.call(vm.global, &.{callback_func});
-        } else {
-            initial_value = callback.call(vm.global, &.{});
+            this.func_arg[this.func_arg.len - 1] = callback_func;
         }
 
+        initial_value = this.func.call(vm.global, @as([]const JSC.JSValue, this.func_arg));
+
         if (initial_value.isAnyError()) {
             if (!Jest.runner.?.did_pending_test_fail) {
                 // test failed unless it's a todo
@@ -730,7 +750,7 @@ pub const TestScope = struct {
             }
         }
 
-        if (callback_length > 0) {
+        if (this.func_has_callback) {
             return .{ .pending = {} };
         }
 
@@ -1034,6 +1054,10 @@ pub const DescribeScope = struct {
         return createScope(globalThis, callframe, "describe.todo()", false, .todo);
     }
 
+    pub fn each(globalThis: *JSGlobalObject, callframe: *CallFrame) callconv(.C) JSValue {
+        return createEach(globalThis, callframe, "describe.each()", "each", false);
+    }
+
     pub fn callIf(globalThis: *JSGlobalObject, callframe: *CallFrame) callconv(.C) JSValue {
         return createIfScope(globalThis, callframe, "describe.if()", "if", DescribeScope, false);
     }
@@ -1042,7 +1066,7 @@ pub const DescribeScope = struct {
         return createIfScope(globalThis, callframe, "describe.skipIf()", "skipIf", DescribeScope, true);
     }
 
-    pub fn run(this: *DescribeScope, globalObject: *JSC.JSGlobalObject, callback: JSC.JSValue) JSC.JSValue {
+    pub fn run(this: *DescribeScope, globalObject: *JSC.JSGlobalObject, callback: JSC.JSValue, args: []const JSC.JSValue) JSC.JSValue {
         if (comptime is_bindgen) return undefined;
         callback.protect();
         defer callback.unprotect();
@@ -1057,7 +1081,7 @@ pub const DescribeScope = struct {
         {
             JSC.markBinding(@src());
             globalObject.clearTerminationException();
-            var result = callback.call(globalObject, &.{});
+            var result = callback.call(globalObject, args);
 
             if (result.asAnyPromise()) |prom| {
                 globalObject.bunVM().waitForPromise(prom);
@@ -1248,7 +1272,7 @@ pub const TestRunnerTask = struct {
         var test_: TestScope = this.describe.tests.items[test_id];
         describe.current_test_id = test_id;
 
-        if (test_.callback == .zero or (describe.is_skip and test_.tag != .only)) {
+        if (test_.func == .zero or (describe.is_skip and test_.tag != .only)) {
             var tag = if (describe.is_skip) describe.tag else test_.tag;
             switch (tag) {
                 .todo => {
@@ -1560,11 +1584,22 @@ inline fn createScope(
             function.protect();
         }
 
+        const func_params_length = function.getLength(globalThis);
+        var arg_size: usize = 0;
+        var has_callback = false;
+        if (func_params_length > 0) {
+            has_callback = true;
+            arg_size = 1;
+        }
+        var function_args = allocator.alloc(JSC.JSValue, arg_size) catch unreachable;
+
         parent.tests.append(allocator, TestScope{
             .label = label,
             .parent = parent,
             .tag = tag_to_use,
-            .callback = if (is_skip) .zero else function,
+            .func = if (is_skip) .zero else function,
+            .func_arg = function_args,
+            .func_has_callback = has_callback,
             .timeout_millis = timeout_ms,
         }) catch unreachable;
 
@@ -1583,7 +1618,7 @@ inline fn createScope(
             .is_skip = is_skip or parent.is_skip,
         };
 
-        return scope.run(globalThis, function);
+        return scope.run(globalThis, function, &.{});
     }
 
     return this;
@@ -1737,3 +1772,187 @@ pub fn printGithubAnnotation(exception: *JSC.ZigException) void {
     Output.printError("\n", .{});
     Output.flush();
 }
+
+pub const EachData = struct { strong: JSC.Strong, is_test: bool };
+
+fn eachBind(
+    globalThis: *JSGlobalObject,
+    callframe: *CallFrame,
+) callconv(.C) JSValue {
+    comptime var signature = "eachBind";
+    const callee = callframe.callee();
+    const arguments = callframe.arguments(3);
+    const args = arguments.ptr[0..arguments.len];
+
+    if (args.len < 2) {
+        globalThis.throwPretty("{s} a description and callback function", .{signature});
+        return .zero;
+    }
+
+    var description = args[0];
+    var function = args[1];
+    var options = if (args.len > 2) args[2] else .zero;
+
+    if (function.isEmptyOrUndefinedOrNull() or !function.isCell() or !function.isCallable(globalThis.vm())) {
+        globalThis.throwPretty("{s} expects a function", .{signature});
+        return .zero;
+    }
+
+    var timeout_ms: u32 = Jest.runner.?.default_timeout_ms;
+    if (options.isNumber()) {
+        timeout_ms = @as(u32, @intCast(@max(args[2].coerce(i32, globalThis), 0)));
+    } else if (options.isObject()) {
+        if (options.get(globalThis, "timeout")) |timeout| {
+            if (!timeout.isNumber()) {
+                globalThis.throwPretty("{s} expects timeout to be a number", .{signature});
+                return .zero;
+            }
+            timeout_ms = @as(u32, @intCast(@max(timeout.coerce(i32, globalThis), 0)));
+        }
+        if (options.get(globalThis, "retry")) |retries| {
+            if (!retries.isNumber()) {
+                globalThis.throwPretty("{s} expects retry to be a number", .{signature});
+                return .zero;
+            }
+            // TODO: retry_count = @intCast(u32, @max(retries.coerce(i32, globalThis), 0));
+        }
+        if (options.get(globalThis, "repeats")) |repeats| {
+            if (!repeats.isNumber()) {
+                globalThis.throwPretty("{s} expects repeats to be a number", .{signature});
+                return .zero;
+            }
+            // TODO: repeat_count = @intCast(u32, @max(repeats.coerce(i32, globalThis), 0));
+        }
+    } else if (!options.isEmptyOrUndefinedOrNull()) {
+        globalThis.throwPretty("{s} expects options to be a number or object", .{signature});
+        return .zero;
+    }
+
+    const parent = DescribeScope.active.?;
+
+    if (JSC.getFunctionData(callee)) |data| {
+        const allocator = getAllocator(globalThis);
+        const each_data = bun.cast(*EachData, data);
+        JSC.setFunctionData(callee, null);
+        const array = each_data.*.strong.get() orelse return .zero;
+        defer {
+            each_data.*.strong.deinit();
+            allocator.destroy(each_data);
+        }
+
+        if (array.isUndefinedOrNull() or !array.jsType().isArray()) {
+            return .zero;
+        }
+
+        var iter = array.arrayIterator(globalThis);
+
+        while (iter.next()) |item| {
+            // TODO: node:util.format() the label
+            const label = if (description.isEmptyOrUndefinedOrNull())
+                ""
+            else
+                (description.toSlice(globalThis, allocator).cloneIfNeeded(allocator) catch unreachable).slice();
+
+            const func_params_length = function.getLength(globalThis);
+            const item_is_array = !item.isEmptyOrUndefinedOrNull() and item.jsType().isArray();
+            var arg_size: usize = 1;
+
+            if (item_is_array) {
+                arg_size = item.getLength(globalThis);
+            }
+
+            // add room for callback function
+            const has_callback_function: bool = (func_params_length > arg_size) and each_data.is_test;
+            if (has_callback_function) {
+                arg_size += 1;
+            }
+
+            var function_args = allocator.alloc(JSC.JSValue, arg_size) catch @panic("can't create function_args");
+            var idx: u32 = 0;
+
+            if (item_is_array) {
+                // Spread array as args
+                var item_iter = item.arrayIterator(globalThis);
+                while (item_iter.next()) |array_item| {
+                    if (array_item == .zero) {
+                        allocator.free(function_args);
+                        break;
+                    }
+                    array_item.protect();
+                    function_args[idx] = array_item;
+                    idx += 1;
+                }
+            } else {
+                item.protect();
+                function_args[0] = item;
+            }
+
+            if (each_data.is_test) {
+                function.protect();
+                parent.tests.append(allocator, TestScope{
+                    .label = label,
+                    .parent = parent,
+                    .tag = parent.tag,
+                    .func = function,
+                    .func_arg = function_args,
+                    .func_has_callback = has_callback_function,
+                    .timeout_millis = timeout_ms,
+                }) catch unreachable;
+
+                if (test_elapsed_timer == null) create_timer: {
+                    var timer = allocator.create(std.time.Timer) catch unreachable;
+                    timer.* = std.time.Timer.start() catch break :create_timer;
+                    test_elapsed_timer = timer;
+                }
+            } else {
+                var scope = allocator.create(DescribeScope) catch unreachable;
+                scope.* = .{
+                    .label = label,
+                    .parent = parent,
+                    .file_id = parent.file_id,
+                    .tag = if (parent.is_skip) parent.tag else .pass,
+                    .is_skip = parent.is_skip,
+                };
+
+                const ret = scope.run(globalThis, function, function_args);
+                _ = ret;
+                allocator.free(function_args);
+            }
+        }
+    }
+
+    return .zero;
+}
+
+inline fn createEach(
+    globalThis: *JSGlobalObject,
+    callframe: *CallFrame,
+    comptime property: string,
+    comptime signature: string,
+    comptime is_test: bool,
+) JSValue {
+    const arguments = callframe.arguments(1);
+    const args = arguments.ptr[0..arguments.len];
+
+    if (args.len == 0) {
+        globalThis.throwPretty("{s} expects an array", .{signature});
+        return .zero;
+    }
+
+    var array = args[0];
+    if (!array.jsType().isArray()) {
+        globalThis.throwPretty("{s} expects an array", .{signature});
+        return .zero;
+    }
+
+    const allocator = getAllocator(globalThis);
+    const name = ZigString.static(property);
+    var strong = JSC.Strong.create(array, globalThis);
+    var each_data = allocator.create(EachData) catch unreachable;
+    each_data.* = EachData{
+        .strong = strong,
+        .is_test = is_test,
+    };
+
+    return JSC.NewFunctionWithData(globalThis, name, 3, eachBind, true, each_data);
+}
diff --git a/test/js/bun/test/jest-each.test.ts b/test/js/bun/test/jest-each.test.ts
new file mode 100644
index 000000000..fc4145cbe
--- /dev/null
+++ b/test/js/bun/test/jest-each.test.ts
@@ -0,0 +1,51 @@
+import { describe, expect, it } from "bun:test";
+
+const NUMBERS = [
+  [1, 1, 2],
+  [1, 2, 3],
+  [2, 1, 3],
+];
+
+describe("jest-each", () => {
+  it("check types", () => {
+    expect(it.each).toBeTypeOf("function");
+    expect(it.each([])).toBeTypeOf("function");
+  });
+  it.each(NUMBERS)("add two numbers", (a, b, e) => {
+    expect(a + b).toBe(e);
+  });
+  it.each(NUMBERS)("add two numbers with callback", (a, b, e, done) => {
+    expect(a + b).toBe(e);
+    expect(done).toBeDefined();
+    // We cast here because we cannot type done when typing args as ...T
+    (done as unknown as (err?: unknown) => void)();
+  });
+  it.each([
+    ["a", "b", "ab"],
+    ["c", "d", "cd"],
+    ["e", "f", "ef"],
+  ])(`adds two strings`, (a, b, res) => {
+    expect(typeof a).toBe("string");
+    expect(typeof b).toBe("string");
+    expect(typeof res).toBe("string");
+    expect(a.concat(b)).toBe(res);
+  });
+  it.each([
+    { a: 1, b: 1, e: 2 },
+    { a: 1, b: 2, e: 3 },
+    { a: 2, b: 13, e: 15 },
+    { a: 2, b: 13, e: 15 },
+    { a: 2, b: 123, e: 125 },
+    { a: 15, b: 13, e: 28 },
+  ])("add two numbers with object", ({ a, b, e }, cb) => {
+    expect(a + b).toBe(e);
+    cb();
+  });
+});
+
+describe.each(["some", "cool", "strings"])("works with describe", s => {
+  it(`has access to params : ${s}`, done => {
+    expect(s).toBeTypeOf("string");
+    done();
+  });
+});