refactor(test): standardize env helpers across suites

src/agents/model-auth.e2e.test.ts

@@ -3,7 +3,7 @@ import os from "node:os";
 import path from "node:path";
 import type { Api, Model } from "@mariozechner/pi-ai";
 import { describe, expect, it } from "vitest";
-import { captureEnv } from "../test-utils/env.js";
+import { withEnvAsync } from "../test-utils/env.js";
 import { ensureAuthProfileStore } from "./auth-profiles.js";
 import { getApiKeyForModel, resolveApiKeyForProvider, resolveEnvApiKey } from "./model-auth.js";
 
@@ -27,38 +27,6 @@ const BEDROCK_PROVIDER_CFG = {
   },
 } as const;
 
-function captureBedrockEnv() {
-  return {
-    bearer: process.env.AWS_BEARER_TOKEN_BEDROCK,
-    access: process.env.AWS_ACCESS_KEY_ID,
-    secret: process.env.AWS_SECRET_ACCESS_KEY,
-    profile: process.env.AWS_PROFILE,
-  };
-}
-
-function restoreBedrockEnv(previous: ReturnType<typeof captureBedrockEnv>) {
-  if (previous.bearer === undefined) {
-    delete process.env.AWS_BEARER_TOKEN_BEDROCK;
-  } else {
-    process.env.AWS_BEARER_TOKEN_BEDROCK = previous.bearer;
-  }
-  if (previous.access === undefined) {
-    delete process.env.AWS_ACCESS_KEY_ID;
-  } else {
-    process.env.AWS_ACCESS_KEY_ID = previous.access;
-  }
-  if (previous.secret === undefined) {
-    delete process.env.AWS_SECRET_ACCESS_KEY;
-  } else {
-    process.env.AWS_SECRET_ACCESS_KEY = previous.secret;
-  }
-  if (previous.profile === undefined) {
-    delete process.env.AWS_PROFILE;
-  } else {
-    process.env.AWS_PROFILE = previous.profile;
-  }
-}
-
 async function resolveBedrockProvider() {
   return resolveApiKeyForProvider({
     provider: "amazon-bedrock",
@@ -67,39 +35,19 @@ async function resolveBedrockProvider() {
   });
 }
 
-async function withEnvUpdates<T>(
-  updates: Record<string, string | undefined>,
-  run: () => Promise<T>,
-): Promise<T> {
-  const snapshot = captureEnv(Object.keys(updates));
-  try {
-    for (const [key, value] of Object.entries(updates)) {
-      if (value === undefined) {
-        delete process.env[key];
-      } else {
-        process.env[key] = value;
-      }
-    }
-    return await run();
-  } finally {
-    snapshot.restore();
-  }
-}
-
 describe("getApiKeyForModel", () => {
   it("migrates legacy oauth.json into auth-profiles.json", async () => {
-    const envSnapshot = captureEnv([
-      "OPENCLAW_STATE_DIR",
-      "OPENCLAW_AGENT_DIR",
-      "PI_CODING_AGENT_DIR",
-    ]);
     const tempDir = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-oauth-"));
 
     try {
-      process.env.OPENCLAW_STATE_DIR = tempDir;
+      const agentDir = path.join(tempDir, "agent");
-      process.env.OPENCLAW_AGENT_DIR = path.join(tempDir, "agent");
+      await withEnvAsync(
-      process.env.PI_CODING_AGENT_DIR = process.env.OPENCLAW_AGENT_DIR;
+        {
-
+          OPENCLAW_STATE_DIR: tempDir,
+          OPENCLAW_AGENT_DIR: agentDir,
+          PI_CODING_AGENT_DIR: agentDir,
+        },
+        async () => {
           const oauthDir = path.join(tempDir, "credentials");
           await fs.mkdir(oauthDir, { recursive: true, mode: 0o700 });
           await fs.writeFile(
@@ -147,27 +95,26 @@ describe("getApiKeyForModel", () => {
               refresh: oauthFixture.refresh,
             },
           });
+        },
+      );
     } finally {
-      envSnapshot.restore();
       await fs.rm(tempDir, { recursive: true, force: true });
     }
   });
 
   it("suggests openai-codex when only Codex OAuth is configured", async () => {
-    const envSnapshot = captureEnv([
-      "OPENAI_API_KEY",
-      "OPENCLAW_STATE_DIR",
-      "OPENCLAW_AGENT_DIR",
-      "PI_CODING_AGENT_DIR",
-    ]);
     const tempDir = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-auth-"));
 
     try {
-      delete process.env.OPENAI_API_KEY;
+      const agentDir = path.join(tempDir, "agent");
-      process.env.OPENCLAW_STATE_DIR = tempDir;
+      await withEnvAsync(
-      process.env.OPENCLAW_AGENT_DIR = path.join(tempDir, "agent");
+        {
-      process.env.PI_CODING_AGENT_DIR = process.env.OPENCLAW_AGENT_DIR;
+          OPENAI_API_KEY: undefined,
-
+          OPENCLAW_STATE_DIR: tempDir,
+          OPENCLAW_AGENT_DIR: agentDir,
+          PI_CODING_AGENT_DIR: agentDir,
+        },
+        async () => {
           const authProfilesPath = path.join(tempDir, "agent", "auth-profiles.json");
           await fs.mkdir(path.dirname(authProfilesPath), {
             recursive: true,
@@ -199,14 +146,15 @@ describe("getApiKeyForModel", () => {
             error = err;
           }
           expect(String(error)).toContain("openai-codex/gpt-5.3-codex");
+        },
+      );
     } finally {
-      envSnapshot.restore();
       await fs.rm(tempDir, { recursive: true, force: true });
     }
   });
 
   it("throws when ZAI API key is missing", async () => {
-    await withEnvUpdates(
+    await withEnvAsync(
       {
         ZAI_API_KEY: undefined,
         Z_AI_API_KEY: undefined,
@@ -228,7 +176,7 @@ describe("getApiKeyForModel", () => {
   });
 
   it("accepts legacy Z_AI_API_KEY for zai", async () => {
-    await withEnvUpdates(
+    await withEnvAsync(
       {
         ZAI_API_KEY: undefined,
         Z_AI_API_KEY: "zai-test-key",
@@ -245,7 +193,7 @@ describe("getApiKeyForModel", () => {
   });
 
   it("resolves Synthetic API key from env", async () => {
-    await withEnvUpdates({ SYNTHETIC_API_KEY: "synthetic-test-key" }, async () => {
+    await withEnvAsync({ SYNTHETIC_API_KEY: "synthetic-test-key" }, async () => {
       const resolved = await resolveApiKeyForProvider({
         provider: "synthetic",
         store: { version: 1, profiles: {} },
@@ -256,7 +204,7 @@ describe("getApiKeyForModel", () => {
   });
 
   it("resolves Qianfan API key from env", async () => {
-    await withEnvUpdates({ QIANFAN_API_KEY: "qianfan-test-key" }, async () => {
+    await withEnvAsync({ QIANFAN_API_KEY: "qianfan-test-key" }, async () => {
       const resolved = await resolveApiKeyForProvider({
         provider: "qianfan",
         store: { version: 1, profiles: {} },
@@ -267,7 +215,7 @@ describe("getApiKeyForModel", () => {
   });
 
   it("resolves Vercel AI Gateway API key from env", async () => {
-    await withEnvUpdates({ AI_GATEWAY_API_KEY: "gateway-test-key" }, async () => {
+    await withEnvAsync({ AI_GATEWAY_API_KEY: "gateway-test-key" }, async () => {
       const resolved = await resolveApiKeyForProvider({
         provider: "vercel-ai-gateway",
         store: { version: 1, profiles: {} },
@@ -278,75 +226,72 @@ describe("getApiKeyForModel", () => {
   });
 
   it("prefers Bedrock bearer token over access keys and profile", async () => {
-    const previous = captureBedrockEnv();
+    await withEnvAsync(
-
+      {
-    try {
+        AWS_BEARER_TOKEN_BEDROCK: "bedrock-token",
-      process.env.AWS_BEARER_TOKEN_BEDROCK = "bedrock-token";
+        AWS_ACCESS_KEY_ID: "access-key",
-      process.env.AWS_ACCESS_KEY_ID = "access-key";
+        AWS_SECRET_ACCESS_KEY: "secret-key",
-      process.env.AWS_SECRET_ACCESS_KEY = "secret-key";
+        AWS_PROFILE: "profile",
-      process.env.AWS_PROFILE = "profile";
+      },
-
+      async () => {
         const resolved = await resolveBedrockProvider();
 
         expect(resolved.mode).toBe("aws-sdk");
         expect(resolved.apiKey).toBeUndefined();
         expect(resolved.source).toContain("AWS_BEARER_TOKEN_BEDROCK");
-    } finally {
+      },
-      restoreBedrockEnv(previous);
+    );
-    }
   });
 
   it("prefers Bedrock access keys over profile", async () => {
-    const previous = captureBedrockEnv();
+    await withEnvAsync(
-
+      {
-    try {
+        AWS_BEARER_TOKEN_BEDROCK: undefined,
-      delete process.env.AWS_BEARER_TOKEN_BEDROCK;
+        AWS_ACCESS_KEY_ID: "access-key",
-      process.env.AWS_ACCESS_KEY_ID = "access-key";
+        AWS_SECRET_ACCESS_KEY: "secret-key",
-      process.env.AWS_SECRET_ACCESS_KEY = "secret-key";
+        AWS_PROFILE: "profile",
-      process.env.AWS_PROFILE = "profile";
+      },
-
+      async () => {
         const resolved = await resolveBedrockProvider();
 
         expect(resolved.mode).toBe("aws-sdk");
         expect(resolved.apiKey).toBeUndefined();
         expect(resolved.source).toContain("AWS_ACCESS_KEY_ID");
-    } finally {
+      },
-      restoreBedrockEnv(previous);
+    );
-    }
   });
 
   it("uses Bedrock profile when access keys are missing", async () => {
-    const previous = captureBedrockEnv();
+    await withEnvAsync(
-
+      {
-    try {
+        AWS_BEARER_TOKEN_BEDROCK: undefined,
-      delete process.env.AWS_BEARER_TOKEN_BEDROCK;
+        AWS_ACCESS_KEY_ID: undefined,
-      delete process.env.AWS_ACCESS_KEY_ID;
+        AWS_SECRET_ACCESS_KEY: undefined,
-      delete process.env.AWS_SECRET_ACCESS_KEY;
+        AWS_PROFILE: "profile",
-      process.env.AWS_PROFILE = "profile";
+      },
-
+      async () => {
         const resolved = await resolveBedrockProvider();
 
         expect(resolved.mode).toBe("aws-sdk");
         expect(resolved.apiKey).toBeUndefined();
         expect(resolved.source).toContain("AWS_PROFILE");
-    } finally {
+      },
-      restoreBedrockEnv(previous);
+    );
-    }
   });
 
   it("accepts VOYAGE_API_KEY for voyage", async () => {
-    await withEnvUpdates({ VOYAGE_API_KEY: "voyage-test-key" }, async () => {
+    await withEnvAsync({ VOYAGE_API_KEY: "voyage-test-key" }, async () => {
-      const resolved = await resolveApiKeyForProvider({
+      const voyage = await resolveApiKeyForProvider({
         provider: "voyage",
         store: { version: 1, profiles: {} },
       });
-      expect(resolved.apiKey).toBe("voyage-test-key");
+      expect(voyage.apiKey).toBe("voyage-test-key");
-      expect(resolved.source).toContain("VOYAGE_API_KEY");
+      expect(voyage.source).toContain("VOYAGE_API_KEY");
     });
   });
 
   it("strips embedded CR/LF from ANTHROPIC_API_KEY", async () => {
-    await withEnvUpdates({ ANTHROPIC_API_KEY: "sk-ant-test-\r\nkey" }, async () => {
+    await withEnvAsync({ ANTHROPIC_API_KEY: "sk-ant-test-\r\nkey" }, async () => {
       const resolved = resolveEnvApiKey("anthropic");
       expect(resolved?.apiKey).toBe("sk-ant-test-key");
       expect(resolved?.source).toContain("ANTHROPIC_API_KEY");
@@ -354,7 +299,7 @@ describe("getApiKeyForModel", () => {
   });
 
   it("resolveEnvApiKey('huggingface') returns HUGGINGFACE_HUB_TOKEN when set", async () => {
-    await withEnvUpdates(
+    await withEnvAsync(
       {
         HUGGINGFACE_HUB_TOKEN: "hf_hub_xyz",
         HF_TOKEN: undefined,
@@ -368,7 +313,7 @@ describe("getApiKeyForModel", () => {
   });
 
   it("resolveEnvApiKey('huggingface') prefers HUGGINGFACE_HUB_TOKEN over HF_TOKEN when both set", async () => {
-    await withEnvUpdates(
+    await withEnvAsync(
       {
         HUGGINGFACE_HUB_TOKEN: "hf_hub_first",
         HF_TOKEN: "hf_second",
@@ -382,7 +327,7 @@ describe("getApiKeyForModel", () => {
   });
 
   it("resolveEnvApiKey('huggingface') returns HF_TOKEN when only HF_TOKEN set", async () => {
-    await withEnvUpdates(
+    await withEnvAsync(
       {
         HUGGINGFACE_HUB_TOKEN: undefined,
         HF_TOKEN: "hf_abc123",

src/browser/config.test.ts

@@ -1,4 +1,5 @@
 import { describe, expect, it } from "vitest";
+import { withEnv } from "../test-utils/env.js";
 import { resolveBrowserConfig, resolveProfile, shouldStartLocalBrowserServer } from "./config.js";
 
 describe("browser config", () => {
@@ -25,9 +26,7 @@ describe("browser config", () => {
   });
 
   it("derives default ports from OPENCLAW_GATEWAY_PORT when unset", () => {
-    const prev = process.env.OPENCLAW_GATEWAY_PORT;
+    withEnv({ OPENCLAW_GATEWAY_PORT: "19001" }, () => {
-    process.env.OPENCLAW_GATEWAY_PORT = "19001";
-    try {
       const resolved = resolveBrowserConfig(undefined);
       expect(resolved.controlPort).toBe(19003);
       const chrome = resolveProfile(resolved, "chrome");
@@ -38,19 +37,11 @@ describe("browser config", () => {
       const openclaw = resolveProfile(resolved, "openclaw");
       expect(openclaw?.cdpPort).toBe(19012);
       expect(openclaw?.cdpUrl).toBe("http://127.0.0.1:19012");
-    } finally {
+    });
-      if (prev === undefined) {
-        delete process.env.OPENCLAW_GATEWAY_PORT;
-      } else {
-        process.env.OPENCLAW_GATEWAY_PORT = prev;
-      }
-    }
   });
 
   it("derives default ports from gateway.port when env is unset", () => {
-    const prev = process.env.OPENCLAW_GATEWAY_PORT;
+    withEnv({ OPENCLAW_GATEWAY_PORT: undefined }, () => {
-    delete process.env.OPENCLAW_GATEWAY_PORT;
-    try {
       const resolved = resolveBrowserConfig(undefined, { gateway: { port: 19011 } });
       expect(resolved.controlPort).toBe(19013);
       const chrome = resolveProfile(resolved, "chrome");
@@ -61,13 +52,7 @@ describe("browser config", () => {
       const openclaw = resolveProfile(resolved, "openclaw");
       expect(openclaw?.cdpPort).toBe(19022);
       expect(openclaw?.cdpUrl).toBe("http://127.0.0.1:19022");
-    } finally {
+    });
-      if (prev === undefined) {
-        delete process.env.OPENCLAW_GATEWAY_PORT;
-      } else {
-        process.env.OPENCLAW_GATEWAY_PORT = prev;
-      }
-    }
   });
 
   it("normalizes hex colors", () => {

src/browser/extension-relay.test.ts

@@ -1,6 +1,7 @@
 import { createServer } from "node:http";
 import { afterEach, beforeEach, describe, expect, it } from "vitest";
 import WebSocket from "ws";
+import { captureEnv } from "../test-utils/env.js";
 import {
   ensureChromeExtensionRelayServer,
   getChromeExtensionRelayAuthHeaders,
@@ -124,10 +125,10 @@ async function waitForListMatch<T>(
 describe("chrome extension relay server", () => {
   const TEST_GATEWAY_TOKEN = "test-gateway-token";
   let cdpUrl = "";
-  let previousGatewayToken: string | undefined;
+  let envSnapshot: ReturnType<typeof captureEnv>;
 
   beforeEach(() => {
-    previousGatewayToken = process.env.OPENCLAW_GATEWAY_TOKEN;
+    envSnapshot = captureEnv(["OPENCLAW_GATEWAY_TOKEN"]);
     process.env.OPENCLAW_GATEWAY_TOKEN = TEST_GATEWAY_TOKEN;
   });
 
@@ -136,11 +137,7 @@ describe("chrome extension relay server", () => {
       await stopChromeExtensionRelayServer({ cdpUrl }).catch(() => {});
       cdpUrl = "";
     }
-    if (previousGatewayToken === undefined) {
+    envSnapshot.restore();
-      delete process.env.OPENCLAW_GATEWAY_TOKEN;
-    } else {
-      process.env.OPENCLAW_GATEWAY_TOKEN = previousGatewayToken;
-    }
   });
 
   it("advertises CDP WS only when extension is connected", async () => {
@@ -438,8 +435,6 @@ describe("chrome extension relay server", () => {
       fakeRelay.once("error", reject);
     });
 
-    const prev = process.env.OPENCLAW_GATEWAY_TOKEN;
-    process.env.OPENCLAW_GATEWAY_TOKEN = "test-gateway-token";
     try {
       cdpUrl = `http://127.0.0.1:${port}`;
       const relay = await ensureChromeExtensionRelayServer({ cdpUrl });
@@ -451,11 +446,6 @@ describe("chrome extension relay server", () => {
       expect(probeToken).toBeTruthy();
       expect(probeToken).not.toBe("test-gateway-token");
     } finally {
-      if (prev === undefined) {
-        delete process.env.OPENCLAW_GATEWAY_TOKEN;
-      } else {
-        process.env.OPENCLAW_GATEWAY_TOKEN = prev;
-      }
       await new Promise<void>((resolve) => fakeRelay.close(() => resolve()));
     }
   });

src/node-host/invoke.sanitize-env.test.ts

@@ -1,31 +1,18 @@
 import { describe, expect, it } from "vitest";
+import { withEnv } from "../test-utils/env.js";
 import { sanitizeEnv } from "./invoke.js";
 import { buildNodeInvokeResultParams } from "./runner.js";
 
 describe("node-host sanitizeEnv", () => {
   it("ignores PATH overrides", () => {
-    const prev = process.env.PATH;
+    withEnv({ PATH: "/usr/bin" }, () => {
-    process.env.PATH = "/usr/bin";
-    try {
       const env = sanitizeEnv({ PATH: "/tmp/evil:/usr/bin" });
       expect(env.PATH).toBe("/usr/bin");
-    } finally {
+    });
-      if (prev === undefined) {
-        delete process.env.PATH;
-      } else {
-        process.env.PATH = prev;
-      }
-    }
   });
 
   it("blocks dangerous env keys/prefixes", () => {
-    const prevPythonPath = process.env.PYTHONPATH;
+    withEnv({ PYTHONPATH: undefined, LD_PRELOAD: undefined, BASH_ENV: undefined }, () => {
-    const prevLdPreload = process.env.LD_PRELOAD;
-    const prevBashEnv = process.env.BASH_ENV;
-    try {
-      delete process.env.PYTHONPATH;
-      delete process.env.LD_PRELOAD;
-      delete process.env.BASH_ENV;
       const env = sanitizeEnv({
         PYTHONPATH: "/tmp/pwn",
         LD_PRELOAD: "/tmp/pwn.so",
@@ -36,46 +23,15 @@ describe("node-host sanitizeEnv", () => {
       expect(env.PYTHONPATH).toBeUndefined();
       expect(env.LD_PRELOAD).toBeUndefined();
       expect(env.BASH_ENV).toBeUndefined();
-    } finally {
+    });
-      if (prevPythonPath === undefined) {
-        delete process.env.PYTHONPATH;
-      } else {
-        process.env.PYTHONPATH = prevPythonPath;
-      }
-      if (prevLdPreload === undefined) {
-        delete process.env.LD_PRELOAD;
-      } else {
-        process.env.LD_PRELOAD = prevLdPreload;
-      }
-      if (prevBashEnv === undefined) {
-        delete process.env.BASH_ENV;
-      } else {
-        process.env.BASH_ENV = prevBashEnv;
-      }
-    }
   });
 
   it("drops dangerous inherited env keys even without overrides", () => {
-    const prevPath = process.env.PATH;
+    withEnv({ PATH: "/usr/bin:/bin", BASH_ENV: "/tmp/pwn.sh" }, () => {
-    const prevBashEnv = process.env.BASH_ENV;
-    try {
-      process.env.PATH = "/usr/bin:/bin";
-      process.env.BASH_ENV = "/tmp/pwn.sh";
       const env = sanitizeEnv(undefined);
       expect(env.PATH).toBe("/usr/bin:/bin");
       expect(env.BASH_ENV).toBeUndefined();
-    } finally {
+    });
-      if (prevPath === undefined) {
-        delete process.env.PATH;
-      } else {
-        process.env.PATH = prevPath;
-      }
-      if (prevBashEnv === undefined) {
-        delete process.env.BASH_ENV;
-      } else {
-        process.env.BASH_ENV = prevBashEnv;
-      }
-    }
   });
 });
 

src/test-utils/env.test.ts

@@ -40,6 +40,20 @@ describe("env test utils", () => {
     expect(process.env[key]).toBe(prev);
   });
 
+  it("withEnv restores values when callback throws", () => {
+    const key = "OPENCLAW_ENV_TEST_SYNC_THROW";
+    const prev = process.env[key];
+
+    expect(() =>
+      withEnv({ [key]: "inside" }, () => {
+        expect(process.env[key]).toBe("inside");
+        throw new Error("boom");
+      }),
+    ).toThrow("boom");
+
+    expect(process.env[key]).toBe(prev);
+  });
+
   it("withEnv can delete a key only inside callback", () => {
     const key = "OPENCLAW_ENV_TEST_SYNC_DELETE";
     const prev = process.env[key];