From 8d4b6ab86f1f75d1132435491523a8d6f319442c Mon Sep 17 00:00:00 2001
From: Michal Piechowiak <misiek.piechowiak@gmail.com>
Date: Wed, 16 Apr 2025 19:38:13 +0200
Subject: [PATCH] test: wait for background work to finish in integration tests

---
 tests/utils/fixture.ts | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/tests/utils/fixture.ts b/tests/utils/fixture.ts
index 59e8b7b09e..3ab24f89c9 100644
--- a/tests/utils/fixture.ts
+++ b/tests/utils/fixture.ts
@@ -396,6 +396,12 @@ export async function invokeFunction(
       process.env[key] = environment[key]
     })
 
+    let resolveInvocation, rejectInvocation
+    const invocationPromise = new Promise((resolve, reject) => {
+      resolveInvocation = resolve
+      rejectInvocation = reject
+    })
+
     const response = (await execute({
       event: {
         headers: headers || {},
@@ -405,8 +411,20 @@ export async function invokeFunction(
       },
       lambdaFunc: { handler },
       timeoutMs: 4_000,
+      onInvocationEnd: (error) => {
+        // lambda-local resolve promise return from execute when response is closed
+        // but we should wait for tracked background work to finish
+        // before resolving the promise to allow background work to finish
+        if (error) {
+          rejectInvocation(error)
+        } else {
+          resolveInvocation()
+        }
+      },
     })) as LambdaResponse
 
+    await invocationPromise
+
     const responseHeaders = Object.entries(response.multiValueHeaders || {}).reduce(
       (prev, [key, value]) => ({
         ...prev,