From fbc54c91518428553aa9444fa4273e7b495ca6be Mon Sep 17 00:00:00 2001
From: Vladimir <sleuths.slews0s@icloud.com>
Date: Mon, 13 Feb 2023 17:37:14 +0100
Subject: [PATCH] fix: wait for console.log to print a message before
 terminating a worker (#2861)

---
 packages/vitest/src/runtime/rpc.ts    | 20 +++++++++++++++++++-
 packages/vitest/src/runtime/worker.ts |  5 +++--
 2 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/packages/vitest/src/runtime/rpc.ts b/packages/vitest/src/runtime/rpc.ts
index 796dfaef40bd..524ac976de00 100644
--- a/packages/vitest/src/runtime/rpc.ts
+++ b/packages/vitest/src/runtime/rpc.ts
@@ -23,12 +23,30 @@ function withSafeTimers(fn: () => void) {
   }
 }
 
+const promises = new Set<Promise<unknown>>()
+
+export const rpcDone = () => {
+  if (!promises.size)
+    return
+  const awaitable = Array.from(promises)
+  return Promise.all(awaitable)
+}
+
 export const rpc = () => {
   const { rpc } = getWorkerState()
   return new Proxy(rpc, {
     get(target, p, handler) {
       const sendCall = Reflect.get(target, p, handler)
-      const safeSendCall = (...args: any[]) => withSafeTimers(() => sendCall(...args))
+      const safeSendCall = (...args: any[]) => withSafeTimers(async () => {
+        const result = sendCall(...args)
+        promises.add(result)
+        try {
+          return await result
+        }
+        finally {
+          promises.delete(result)
+        }
+      })
       safeSendCall.asEvent = sendCall.asEvent
       return safeSendCall
     },
diff --git a/packages/vitest/src/runtime/worker.ts b/packages/vitest/src/runtime/worker.ts
index eb19e5e1564e..e7528b82b416 100644
--- a/packages/vitest/src/runtime/worker.ts
+++ b/packages/vitest/src/runtime/worker.ts
@@ -11,7 +11,7 @@ import { getWorkerState } from '../utils/global'
 import type { MockMap } from '../types/mocker'
 import type { VitestExecutor } from './execute'
 import { createVitestExecutor } from './execute'
-import { rpc } from './rpc'
+import { rpc, rpcDone } from './rpc'
 
 let _viteNode: {
   run: (files: string[], config: ResolvedConfig, executor: VitestExecutor) => Promise<void>
@@ -109,5 +109,6 @@ function init(ctx: WorkerContext) {
 export async function run(ctx: WorkerContext) {
   init(ctx)
   const { run, executor } = await startViteNode(ctx)
-  return run(ctx.files, ctx.config, executor)
+  await run(ctx.files, ctx.config, executor)
+  await rpcDone()
 }