GitklcHub
diff --git a/‎heft-plugins/heft-jest-plugin/includes/jestScriptTransformerPatch.js
Lines changed: 29 additions & 7 deletions b/‎heft-plugins/heft-jest-plugin/includes/jestScriptTransformerPatch.js
Lines changed: 29 additions & 7 deletions
@@ -16,11 +16,6 @@ const patchName = path.basename(__filename);
 const HEFT_JEST_DISABLE_CACHE_ENV_VARIABLE = 'HEFT_JEST_DISABLE_CACHE';
 
 function applyPatch() {
-  // Only apply the patch if the environment variable is set
-  if (!process.env[HEFT_JEST_DISABLE_CACHE_ENV_VARIABLE]) {
-    return;
-  }
-
   try {
     let contextFolder = __dirname;
     // Resolve the "@jest/core" package relative to heft-jest-plugin
@@ -54,7 +49,20 @@ function patchScriptTransformer(scriptPath) {
   const functionsToReplace = ['readCacheFile', 'writeCacheFile'];
 
   for (const functionName of functionsToReplace) {
-    const match = scriptContent.match(new RegExp(`^\\s*const ${functionName} =`, 'm'));
+    const originalFunctionName = `${functionName}Original`;
+
+    // First, attempt to remove the existing patch, if one is present
+    let match = scriptContent.match(new RegExp(`^\\s*const ${originalFunctionName} =`, 'm'));
+    let originalFunctionContent;
+    if (match) {
+      const startIndex = match.index;
+      const endIndex = scriptContent.indexOf('};', startIndex) + 2;
+      originalFunctionContent = scriptContent.slice(startIndex, endIndex);
+      scriptContent = scriptContent.slice(0, startIndex) + scriptContent.slice(endIndex);
+    }
+
+    // Then, attempt to find the function to patch
+    match = scriptContent.match(new RegExp(`^\\s*const ${functionName} =`, 'm'));
     if (!match) {
       throw new Error(
         `The ${JSON.stringify(functionName)} function was not found in the file ${JSON.stringify(scriptPath)}`
@@ -63,9 +71,23 @@ function patchScriptTransformer(scriptPath) {
 
     const startIndex = match.index;
     const endIndex = scriptContent.indexOf('};', startIndex) + 2;
+
+    // If we already have the original function content, no need to extract it again
+    if (originalFunctionContent === undefined) {
+      originalFunctionContent = scriptContent
+        .slice(startIndex, endIndex)
+        .replace(`const ${functionName} =`, `const ${originalFunctionName} =`);
+    }
+
     scriptContent =
       scriptContent.slice(0, startIndex) +
-      `const ${functionName} = () => {};` +
+      `${originalFunctionContent}\n` +
+      `const ${functionName} = (...args) => {\n` +
+      `  if (process.env['${HEFT_JEST_DISABLE_CACHE_ENV_VARIABLE}']) {\n` +
+      `    return;\n` +
+      `  }\n` +
+      `  return ${originalFunctionName}(...args);\n` +
+      '};' +
       scriptContent.slice(endIndex);
   }