refactor(test): share env var helpers

src/agents/tools/web-search.e2e.test.ts

@@ -1,30 +1,7 @@
 import { describe, expect, it } from "vitest";
+import { withEnv } from "../../test-utils/env.js";
 import { __testing } from "./web-search.js";
 
-function withEnv<T>(env: Record<string, string | undefined>, fn: () => T): T {
-  const prev: Record<string, string | undefined> = {};
-  for (const [key, value] of Object.entries(env)) {
-    prev[key] = process.env[key];
-    if (value === undefined) {
-      // Make tests hermetic even on machines with real keys set.
-      delete process.env[key];
-    } else {
-      process.env[key] = value;
-    }
-  }
-  try {
-    return fn();
-  } finally {
-    for (const [key, value] of Object.entries(prev)) {
-      if (value === undefined) {
-        delete process.env[key];
-      } else {
-        process.env[key] = value;
-      }
-    }
-  }
-}
-
 const {
   inferPerplexityBaseUrlFromApiKey,
   resolvePerplexityBaseUrl,

src/commands/onboard-non-interactive.provider-auth.e2e.test.ts

@@ -3,6 +3,7 @@ import os from "node:os";
 import path from "node:path";
 import { setTimeout as delay } from "node:timers/promises";
 import { describe, expect, it } from "vitest";
+import { captureEnv } from "../test-utils/env.js";
 import { MINIMAX_API_BASE_URL, MINIMAX_CN_API_BASE_URL } from "./onboard-auth.js";
 import { OPENAI_DEFAULT_MODEL } from "./openai-model-default.js";
 
@@ -12,20 +13,6 @@ type RuntimeMock = {
   exit: (code: number) => never;
 };
 
-type EnvSnapshot = {
-  home: string | undefined;
-  stateDir: string | undefined;
-  configPath: string | undefined;
-  skipChannels: string | undefined;
-  skipGmail: string | undefined;
-  skipCron: string | undefined;
-  skipCanvas: string | undefined;
-  token: string | undefined;
-  password: string | undefined;
-  customApiKey: string | undefined;
-  disableConfigCache: string | undefined;
-};
-
 type OnboardEnv = {
   configPath: string;
   runtime: RuntimeMock;
@@ -47,49 +34,23 @@ async function removeDirWithRetry(dir: string): Promise<void> {
   }
 }
 
-function captureEnv(): EnvSnapshot {
-  return {
-    home: process.env.HOME,
-    stateDir: process.env.OPENCLAW_STATE_DIR,
-    configPath: process.env.OPENCLAW_CONFIG_PATH,
-    skipChannels: process.env.OPENCLAW_SKIP_CHANNELS,
-    skipGmail: process.env.OPENCLAW_SKIP_GMAIL_WATCHER,
-    skipCron: process.env.OPENCLAW_SKIP_CRON,
-    skipCanvas: process.env.OPENCLAW_SKIP_CANVAS_HOST,
-    token: process.env.OPENCLAW_GATEWAY_TOKEN,
-    password: process.env.OPENCLAW_GATEWAY_PASSWORD,
-    customApiKey: process.env.CUSTOM_API_KEY,
-    disableConfigCache: process.env.OPENCLAW_DISABLE_CONFIG_CACHE,
-  };
-}
-
-function restoreEnvVar(key: keyof NodeJS.ProcessEnv, value: string | undefined): void {
-  if (value == null) {
-    delete process.env[key];
-    return;
-  }
-  process.env[key] = value;
-}
-
-function restoreEnv(prev: EnvSnapshot): void {
-  restoreEnvVar("HOME", prev.home);
-  restoreEnvVar("OPENCLAW_STATE_DIR", prev.stateDir);
-  restoreEnvVar("OPENCLAW_CONFIG_PATH", prev.configPath);
-  restoreEnvVar("OPENCLAW_SKIP_CHANNELS", prev.skipChannels);
-  restoreEnvVar("OPENCLAW_SKIP_GMAIL_WATCHER", prev.skipGmail);
-  restoreEnvVar("OPENCLAW_SKIP_CRON", prev.skipCron);
-  restoreEnvVar("OPENCLAW_SKIP_CANVAS_HOST", prev.skipCanvas);
-  restoreEnvVar("OPENCLAW_GATEWAY_TOKEN", prev.token);
-  restoreEnvVar("OPENCLAW_GATEWAY_PASSWORD", prev.password);
-  restoreEnvVar("CUSTOM_API_KEY", prev.customApiKey);
-  restoreEnvVar("OPENCLAW_DISABLE_CONFIG_CACHE", prev.disableConfigCache);
-}
-
 async function withOnboardEnv(
   prefix: string,
   run: (ctx: OnboardEnv) => Promise<void>,
 ): Promise<void> {
-  const prev = captureEnv();
+  const prev = captureEnv([
+    "HOME",
+    "OPENCLAW_STATE_DIR",
+    "OPENCLAW_CONFIG_PATH",
+    "OPENCLAW_SKIP_CHANNELS",
+    "OPENCLAW_SKIP_GMAIL_WATCHER",
+    "OPENCLAW_SKIP_CRON",
+    "OPENCLAW_SKIP_CANVAS_HOST",
+    "OPENCLAW_GATEWAY_TOKEN",
+    "OPENCLAW_GATEWAY_PASSWORD",
+    "CUSTOM_API_KEY",
+    "OPENCLAW_DISABLE_CONFIG_CACHE",
+  ]);
 
   process.env.OPENCLAW_SKIP_CHANNELS = "1";
   process.env.OPENCLAW_SKIP_GMAIL_WATCHER = "1";
@@ -120,7 +81,7 @@ async function withOnboardEnv(
     await run({ configPath, runtime });
   } finally {
     await removeDirWithRetry(tempHome);
-    restoreEnv(prev);
+    prev.restore();
   }
 }
 

src/test-utils/env.ts

@@ -16,3 +16,38 @@ export function captureEnv(keys: string[]) {
     },
   };
 }
+
+export function withEnv<T>(env: Record<string, string | undefined>, fn: () => T): T {
+  const snapshot = captureEnv(Object.keys(env));
+  try {
+    for (const [key, value] of Object.entries(env)) {
+      if (value === undefined) {
+        delete process.env[key];
+      } else {
+        process.env[key] = value;
+      }
+    }
+    return fn();
+  } finally {
+    snapshot.restore();
+  }
+}
+
+export async function withEnvAsync<T>(
+  env: Record<string, string | undefined>,
+  fn: () => Promise<T>,
+): Promise<T> {
+  const snapshot = captureEnv(Object.keys(env));
+  try {
+    for (const [key, value] of Object.entries(env)) {
+      if (value === undefined) {
+        delete process.env[key];
+      } else {
+        process.env[key] = value;
+      }
+    }
+    return await fn();
+  } finally {
+    snapshot.restore();
+  }
+}

src/tts/tts.test.ts

@@ -2,6 +2,7 @@ import { completeSimple } from "@mariozechner/pi-ai";
 import { describe, expect, it, vi, beforeEach } from "vitest";
 import { getApiKeyForModel } from "../agents/model-auth.js";
 import { resolveModel } from "../agents/pi-embedded-runner/model.js";
+import { withEnv } from "../test-utils/env.js";
 import * as tts from "./tts.js";
 
 vi.mock("@mariozechner/pi-ai", () => ({
@@ -367,38 +368,6 @@ describe("tts", () => {
       messages: { tts: {} },
     };
 
-    const restoreEnv = (snapshot: Record<string, string | undefined>) => {
-      const keys = ["OPENAI_API_KEY", "ELEVENLABS_API_KEY", "XI_API_KEY"] as const;
-      for (const key of keys) {
-        const value = snapshot[key];
-        if (value === undefined) {
-          delete process.env[key];
-        } else {
-          process.env[key] = value;
-        }
-      }
-    };
-
-    const withEnv = (env: Record<string, string | undefined>, run: () => void) => {
-      const snapshot = {
-        OPENAI_API_KEY: process.env.OPENAI_API_KEY,
-        ELEVENLABS_API_KEY: process.env.ELEVENLABS_API_KEY,
-        XI_API_KEY: process.env.XI_API_KEY,
-      };
-      try {
-        for (const [key, value] of Object.entries(env)) {
-          if (value === undefined) {
-            delete process.env[key];
-          } else {
-            process.env[key] = value;
-          }
-        }
-        run();
-      } finally {
-        restoreEnv(snapshot);
-      }
-    };
-
     it("prefers OpenAI when no provider is configured and API key exists", () => {
       withEnv(
         {