From bd9771be7171fad4d32ed6df294ab452db71804c Mon Sep 17 00:00:00 2001
From: Michal Piechowiak <misiek.piechowiak@gmail.com>
Date: Mon, 17 Jun 2024 13:13:02 +0200
Subject: [PATCH] test: don't destroy lambda env vars until response stream
 finished

---
 tests/utils/fixture.ts        | 22 ++++++++++++++++++++--
 tests/utils/sandbox-child.mjs | 21 +++++++++++++++++++--
 2 files changed, 39 insertions(+), 4 deletions(-)

diff --git a/tests/utils/fixture.ts b/tests/utils/fixture.ts
index 44931bec94..7063af22eb 100644
--- a/tests/utils/fixture.ts
+++ b/tests/utils/fixture.ts
@@ -364,14 +364,24 @@ export async function invokeFunction(
       NETLIFY_BLOBS_CONTEXT: createBlobContext(ctx),
       ...(env || {}),
     }
+
+    const envVarsToRestore = {}
+
+    // We are not using lambda-local's environment variable setting because it cleans up
+    // environment vars to early (before stream is closed)
+    Object.keys(environment).forEach(function (key) {
+      if (typeof process.env[key] !== 'undefined') {
+        envVarsToRestore[key] = process.env[key]
+      }
+      process.env[key] = environment[key]
+    })
+
     const response = (await execute({
       event: {
         headers: headers || {},
         httpMethod: httpMethod || 'GET',
         rawUrl: new URL(url || '/', 'https://example.netlify').href,
       },
-      environment,
-      envdestroy: true,
       lambdaFunc: { handler },
       timeoutMs: 4_000,
     })) as LambdaResponse
@@ -386,6 +396,14 @@ export async function invokeFunction(
 
     const bodyBuffer = await streamToBuffer(response.body)
 
+    Object.keys(environment).forEach(function (key) {
+      if (typeof envVarsToRestore[key] !== 'undefined') {
+        process.env[key] = envVarsToRestore[key]
+      } else {
+        delete process.env[key]
+      }
+    })
+
     return {
       statusCode: response.statusCode,
       bodyBuffer,
diff --git a/tests/utils/sandbox-child.mjs b/tests/utils/sandbox-child.mjs
index a6371cbcd0..4439a667c3 100644
--- a/tests/utils/sandbox-child.mjs
+++ b/tests/utils/sandbox-child.mjs
@@ -48,14 +48,23 @@ process.on('message', async (msg) => {
         ...(env || {}),
       }
 
+      const envVarsToRestore = {}
+
+      // We are not using lambda-local's environment variable setting because it cleans up
+      // environment vars to early (before stream is closed)
+      Object.keys(environment).forEach(function (key) {
+        if (typeof process.env[key] !== 'undefined') {
+          envVarsToRestore[key] = process.env[key]
+        }
+        process.env[key] = environment[key]
+      })
+
       const response = await execute({
         event: {
           headers: headers || {},
           httpMethod: httpMethod || 'GET',
           rawUrl: new URL(url || '/', 'https://example.netlify').href,
         },
-        environment,
-        envdestroy: true,
         lambdaFunc: { handler },
         timeoutMs: 4_000,
       })
@@ -70,6 +79,14 @@ process.on('message', async (msg) => {
 
       const bodyBuffer = await streamToBuffer(response.body)
 
+      Object.keys(environment).forEach(function (key) {
+        if (typeof envVarsToRestore[key] !== 'undefined') {
+          process.env[key] = envVarsToRestore[key]
+        } else {
+          delete process.env[key]
+        }
+      })
+
       const result = {
         statusCode: response.statusCode,
         bodyBuffer,