From f43bfccf66cebde336db9051ad40c5029d557f82 Mon Sep 17 00:00:00 2001
From: sunrabbit123 <qudwls185@naver.com>
Date: Mon, 30 Jun 2025 13:33:02 +0900
Subject: [PATCH] test: fix test code for save test case principal

Signed-off-by: sunrabbit123 <qudwls185@naver.com>
---
 src/client/stdio.test.ts | 28 +++++++++++++++++++---------
 1 file changed, 19 insertions(+), 9 deletions(-)

diff --git a/src/client/stdio.test.ts b/src/client/stdio.test.ts
index cc3731fb..eb7f0e1b 100644
--- a/src/client/stdio.test.ts
+++ b/src/client/stdio.test.ts
@@ -1,17 +1,20 @@
 import { JSONRPCMessage } from "../types.js";
 import { StdioClientTransport, StdioServerParameters, DEFAULT_INHERITED_ENV_VARS, getDefaultEnvironment } from "./stdio.js";
+import { AsyncLocalStorage } from "node:async_hooks";
 
 const serverParameters: StdioServerParameters = {
   command: "/usr/bin/tee",
 };
 
-
-let spawnEnv: Record<string, string> | undefined;
+const envAsyncLocalStorage = new AsyncLocalStorage<{ env: Record<string, string> }>();
 
 jest.mock('cross-spawn', () => {
   const originalSpawn = jest.requireActual('cross-spawn');
   return jest.fn((command, args, options) => {
-    spawnEnv = options.env;
+    const env = envAsyncLocalStorage.getStore();
+    if (env) {
+      env.env = options.env;
+    }
     return originalSpawn(command, args, options);
   });
 });
@@ -72,6 +75,7 @@ test("should read messages", async () => {
 });
 
 test("should properly set default environment variables in spawned process", async () => {
+  await envAsyncLocalStorage.run({ env: {} }, async () => {
   const client = new StdioClientTransport(serverParameters);
 
   await client.start();
@@ -79,18 +83,21 @@ test("should properly set default environment variables in spawned process", asy
 
   // Get the default environment variables
   const defaultEnv = getDefaultEnvironment();
-
+  const spawnEnv = envAsyncLocalStorage.getStore()?.env;
+  expect(spawnEnv).toBeDefined();
   // Verify that all default environment variables are present
   for (const key of DEFAULT_INHERITED_ENV_VARS) {
     if (process.env[key] && !process.env[key].startsWith("()")) {
       expect(spawnEnv).toHaveProperty(key);
       expect(spawnEnv![key]).toBe(process.env[key]);
-      expect(spawnEnv![key]).toBe(defaultEnv[key]);
+        expect(spawnEnv![key]).toBe(defaultEnv[key]);
+      }
     }
-  }
+  });
 });
 
 test("should override default environment variables with custom ones", async () => {
+  await envAsyncLocalStorage.run({ env: {} }, async () => {
   const customEnv = {
     HOME: "/custom/home",
     PATH: "/custom/path",
@@ -104,7 +111,9 @@ test("should override default environment variables with custom ones", async ()
 
   await client.start();
   await client.close();
-
+  
+  const spawnEnv = envAsyncLocalStorage.getStore()?.env;
+  expect(spawnEnv).toBeDefined();
   // Verify that custom environment variables override default ones
   for (const [key, value] of Object.entries(customEnv)) {
     expect(spawnEnv).toHaveProperty(key);
@@ -117,5 +126,6 @@ test("should override default environment variables with custom ones", async ()
       expect(spawnEnv).toHaveProperty(key);
       expect(spawnEnv![key]).toBe(process.env[key]);
     }
-  }
-}); 
+    }
+  });
+});