From be9a1f8dc72d16b1845f139f3a40e6a38649ac8f Mon Sep 17 00:00:00 2001
From: Patrik Oldsberg <poldsberg@gmail.com>
Date: Thu, 26 Dec 2024 11:31:03 +0100
Subject: [PATCH] cli: simplify Jest ESM support

Signed-off-by: Patrik Oldsberg <poldsberg@gmail.com>
---
 .../cli/config/jestCachingModuleLoader.js     | 20 ++++++++-----------
 packages/cli/config/jestSwcTransform.js       |  9 +++++++--
 2 files changed, 15 insertions(+), 14 deletions(-)

diff --git a/packages/cli/config/jestCachingModuleLoader.js b/packages/cli/config/jestCachingModuleLoader.js
index 3b7740a4f7..dd22f25f6f 100644
--- a/packages/cli/config/jestCachingModuleLoader.js
+++ b/packages/cli/config/jestCachingModuleLoader.js
@@ -19,6 +19,11 @@ const { default: JestRuntime } = require('jest-runtime');
 const scriptTransformCache = new Map();
 
 module.exports = class CachingJestRuntime extends JestRuntime {
+  constructor(config, ...restAgs) {
+    super(config, ...restAgs);
+    this.allowLoadAsEsm = config.extensionsToTreatAsEsm.includes('.mts');
+  }
+
   // This may or may not be a good idea. Theoretically I don't know why this would impact
   // test correctness and flakiness, but it seems like it may introduce flakiness and strange failures.
   // It does seem to speed up test execution by a fair amount though.
@@ -34,19 +39,10 @@ module.exports = class CachingJestRuntime extends JestRuntime {
     return script;
   }
 
-  // Notes(Rugvip): As far as I can tell this is the best we can currently do
-  // for runtime ESM support in Jest. What the below logic effectively does is
-  // to only allow packages to be loaded as ESM if all imports of that package
-  // are done in an ESM compatible way, as in either from ESM code or with a
-  // dynamic import.
-  cjsModules = new Set();
-  _resolveCjsModule(...args) {
-    const path = super._resolveCjsModule(...args);
-    this.cjsModules.add(path);
-    return path;
-  }
+  // Unfortunately we need to use this unstable API to make sure that .js files
+  // are only loaded as modules where ESM is supported, i.e. Node.js packages.
   unstable_shouldLoadAsEsm(path, ...restArgs) {
-    if (this.cjsModules.has(path)) {
+    if (!this.allowLoadAsEsm) {
       return false;
     }
     return super.unstable_shouldLoadAsEsm(path, ...restArgs);
diff --git a/packages/cli/config/jestSwcTransform.js b/packages/cli/config/jestSwcTransform.js
index 203bb73d80..a77683b939 100644
--- a/packages/cli/config/jestSwcTransform.js
+++ b/packages/cli/config/jestSwcTransform.js
@@ -18,13 +18,18 @@ const { createTransformer: createSwcTransformer } = require('@swc/jest');
 const ESM_REGEX = /\b(?:import|export)\b/;
 
 function createTransformer(config) {
-  const useModules = Boolean(config?.module);
   const swcTransformer = createSwcTransformer({
     inputSourceMap: false,
     ...config,
   });
   const process = (source, filePath, jestOptions) => {
-    if (filePath.endsWith('.js') && (useModules || !ESM_REGEX.test(source))) {
+    // Skip transformation of .js files without ESM syntax, we never transform from CJS to ESM
+    if (filePath.endsWith('.js') && !ESM_REGEX.test(source)) {
+      return { code: source };
+    }
+
+    // Skip transformation of .mjs files, they should only be used if ESM support is available
+    if (filePath.endsWith('.mjs')) {
       return { code: source };
     }