[bun:test] Auto-import jest globals in test files

7 files changed, 99 insertions, 0 deletions

diff --git a/src/bun.js/javascript.zig b/src/bun.js/javascript.zig
index af8f9c501..7aa74f51f 100644
--- a/src/bun.js/javascript.zig
+++ b/src/bun.js/javascript.zig
@@ -350,6 +350,7 @@ pub const VirtualMachine = struct {
     log: *logger.Log,
     event_listeners: EventListenerMixin.Map,
     main: string = "",
+    main_hash: u32 = 0,
     process: js.JSObjectRef = null,
     blobs: ?*Blob.Group = null,
     flush_list: std.ArrayList(string),
diff --git a/src/bun.js/module_loader.zig b/src/bun.js/module_loader.zig
index 5945721b5..dd4e258e6 100644
--- a/src/bun.js/module_loader.zig
+++ b/src/bun.js/module_loader.zig
@@ -964,6 +964,10 @@ pub const ModuleLoader = struct {
                     .jsx = jsc_vm.bundler.options.jsx,
                     .virtual_source = virtual_source,
                     .hoist_bun_plugin = true,
+                    .inject_jest_globals = jsc_vm.bundler.options.rewrite_jest_for_tests and
+                        jsc_vm.main.len == path.text.len and
+                        jsc_vm.main_hash == hash and
+                        strings.eqlLong(jsc_vm.main, path.text, false),
                 };
 
                 if (is_node_override) {
diff --git a/src/bundler.zig b/src/bundler.zig
index 336723da3..f540436b0 100644
--- a/src/bundler.zig
+++ b/src/bundler.zig
@@ -1324,6 +1324,7 @@ pub const Bundler = struct {
         virtual_source: ?*const logger.Source = null,
         replace_exports: runtime.Runtime.Features.ReplaceableExport.Map = .{},
         hoist_bun_plugin: bool = false,
+        inject_jest_globals: bool = false,
     };
 
     pub fn parse(
@@ -1453,6 +1454,7 @@ pub const Bundler = struct {
 
                 opts.features.jsx_optimization_hoist = bundler.options.jsx_optimization_hoist orelse opts.features.jsx_optimization_inline;
                 opts.features.hoist_bun_plugin = this_parse.hoist_bun_plugin;
+                opts.features.inject_jest_globals = this_parse.inject_jest_globals;
                 if (bundler.macro_context == null) {
                     bundler.macro_context = js_ast.Macro.MacroContext.init(bundler);
                 }
diff --git a/src/cli/test_command.zig b/src/cli/test_command.zig
index bc8d7d06a..2b0f82bbd 100644
--- a/src/cli/test_command.zig
+++ b/src/cli/test_command.zig
@@ -561,6 +561,7 @@ pub const TestCommand = struct {
         Output.prettyErrorln("<r>\n{s}:\n", .{resolution.path_pair.primary.name.filename});
         Output.flush();
 
+        vm.main_hash = @truncate(u32, bun.hash(resolution.path_pair.primary.text));
         var promise = try vm.loadEntryPoint(resolution.path_pair.primary.text);
 
         switch (promise.status(vm.global.vm())) {
diff --git a/src/js_ast.zig b/src/js_ast.zig
index 5b575eda5..a3225ccce 100644
--- a/src/js_ast.zig
+++ b/src/js_ast.zig
@@ -4716,6 +4716,7 @@ pub const Part = struct {
         react_fast_refresh,
         dirname_filename,
         bun_plugin,
+        bun_test,
     };
 
     pub const SymbolUseMap = std.ArrayHashMapUnmanaged(Ref, Symbol.Use, RefHashCtx, false);
diff --git a/src/js_parser.zig b/src/js_parser.zig
index 4159faada..27c349f24 100644
--- a/src/js_parser.zig
+++ b/src/js_parser.zig
@@ -2563,6 +2563,70 @@ pub const Parser = struct {
             exports_kind = .esm;
         }
 
+        // Auto inject jest globals into the test file
+        if (p.options.features.inject_jest_globals) outer: {
+            var jest: *Jest = &p.jest;
+
+            for (p.import_records.items) |*item| {
+                // skip if they did import it
+                if (strings.eqlComptime(item.path.text, "bun:test") or strings.eqlComptime(item.path.text, "@jest/globals") or strings.eqlComptime(item.path.text, "vitest")) {
+                    break :outer;
+                }
+            }
+
+            // if they didn't use any of the jest globals, don't inject it, I guess.
+            const items_count = brk: {
+                var count: usize = 0;
+                inline for (comptime std.meta.fieldNames(Jest)) |symbol_name| {
+                    count += @boolToInt(p.symbols.items[@field(jest, symbol_name).innerIndex()].use_count_estimate > 0);
+                }
+
+                break :brk count;
+            };
+            if (items_count == 0)
+                break :outer;
+
+            const import_record_id = p.addImportRecord(.stmt, logger.Loc.Empty, "bun:test");
+            var import_record: *ImportRecord = &p.import_records.items[import_record_id];
+            import_record.tag = .bun_test;
+
+            var declared_symbols = try p.allocator.alloc(js_ast.DeclaredSymbol, items_count);
+            var clauses: []js_ast.ClauseItem = p.allocator.alloc(js_ast.ClauseItem, items_count) catch unreachable;
+            var clause_i: usize = 0;
+            inline for (comptime std.meta.fieldNames(Jest)) |symbol_name| {
+                if (p.symbols.items[@field(jest, symbol_name).innerIndex()].use_count_estimate > 0) {
+                    clauses[clause_i] = js_ast.ClauseItem{
+                        .name = .{ .ref = @field(jest, symbol_name), .loc = logger.Loc.Empty },
+                        .alias = symbol_name,
+                        .alias_loc = logger.Loc.Empty,
+                        .original_name = "",
+                    };
+                    declared_symbols[clause_i] = .{ .ref = @field(jest, symbol_name), .is_top_level = true };
+                    clause_i += 1;
+                }
+            }
+
+            const import_stmt = p.s(
+                S.Import{
+                    .namespace_ref = p.declareSymbol(.unbound, logger.Loc.Empty, "bun_test_import_namespace_for_internal_use_only") catch unreachable,
+                    .items = clauses,
+                    .import_record_index = import_record_id,
+                },
+                logger.Loc.Empty,
+            );
+
+            var part_stmts = try p.allocator.alloc(Stmt, 1);
+            part_stmts[0] = import_stmt;
+            var import_record_indices = try p.allocator.alloc(u32, 1);
+            import_record_indices[0] = import_record_id;
+            before.append(js_ast.Part{
+                .stmts = part_stmts,
+                .declared_symbols = declared_symbols,
+                .import_record_indices = import_record_indices,
+                .tag = .bun_test,
+            }) catch unreachable;
+        }
+
         // Auto-import & post-process JSX
         switch (comptime ParserType.jsx_transform_type) {
             .react => {
@@ -3515,6 +3579,17 @@ pub const MacroState = struct {
     }
 };
 
+const Jest = struct {
+    expect: Ref = Ref.None,
+    describe: Ref = Ref.None,
+    @"test": Ref = Ref.None,
+    it: Ref = Ref.None,
+    beforeEach: Ref = Ref.None,
+    afterEach: Ref = Ref.None,
+    beforeAll: Ref = Ref.None,
+    afterAll: Ref = Ref.None,
+};
+
 // workaround for https://github.com/ziglang/zig/issues/10903
 fn NewParser(
     comptime parser_features: ParserFeatures,
@@ -3656,6 +3731,8 @@ fn NewParser_(
 
         bun_jsx_ref: Ref = Ref.None,
 
+        jest: Jest = .{},
+
         // Imports (both ES6 and CommonJS) are tracked at the top level
         import_records: ImportRecordList,
         import_records_for_current_part: List(u32) = .{},
@@ -5107,6 +5184,17 @@ fn NewParser_(
             p.dirname_ref = try p.declareCommonJSSymbol(.unbound, "__dirname");
             p.filename_ref = try p.declareCommonJSSymbol(.unbound, "__filename");
 
+            if (p.options.features.inject_jest_globals) {
+                p.jest.describe = try p.declareCommonJSSymbol(.unbound, "describe");
+                p.jest.@"test" = try p.declareCommonJSSymbol(.unbound, "test");
+                p.jest.it = try p.declareCommonJSSymbol(.unbound, "it");
+                p.jest.expect = try p.declareCommonJSSymbol(.unbound, "expect");
+                p.jest.beforeEach = try p.declareCommonJSSymbol(.unbound, "beforeEach");
+                p.jest.afterEach = try p.declareCommonJSSymbol(.unbound, "afterEach");
+                p.jest.beforeAll = try p.declareCommonJSSymbol(.unbound, "beforeAll");
+                p.jest.afterAll = try p.declareCommonJSSymbol(.unbound, "afterAll");
+            }
+
             if (p.options.enable_bundling) {
                 p.runtime_imports.__reExport = try p.declareGeneratedSymbol(.other, "__reExport");
                 p.runtime_imports.@"$$m" = try p.declareGeneratedSymbol(.other, "$$m");
diff --git a/src/runtime.zig b/src/runtime.zig
index e72165653..2e9e06dcf 100644
--- a/src/runtime.zig
+++ b/src/runtime.zig
@@ -292,6 +292,8 @@ pub const Runtime = struct {
         allow_runtime: bool = true,
         inlining: bool = false,
 
+        inject_jest_globals: bool = false,
+
         /// Instead of jsx("div", {}, void 0)
         /// ->
         /// {