refactor(test): simplify state dir env helpers

src/agents/auth-profiles/oauth.fallback-to-main-agent.e2e.test.ts

@@ -3,13 +3,16 @@ import os from "node:os";
 import path from "node:path";
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 import type { AuthProfileStore } from "./types.js";
+import { captureEnv } from "../../test-utils/env.js";
 import { resolveApiKeyForProfile } from "./oauth.js";
 import { ensureAuthProfileStore } from "./store.js";
 
 describe("resolveApiKeyForProfile fallback to main agent", () => {
-  const previousStateDir = process.env.OPENCLAW_STATE_DIR;
-  const previousAgentDir = process.env.OPENCLAW_AGENT_DIR;
-  const previousPiAgentDir = process.env.PI_CODING_AGENT_DIR;
+  const envSnapshot = captureEnv([
+    "OPENCLAW_STATE_DIR",
+    "OPENCLAW_AGENT_DIR",
+    "PI_CODING_AGENT_DIR",
+  ]);
   let tmpDir: string;
   let mainAgentDir: string;
   let secondaryAgentDir: string;
@@ -30,22 +33,7 @@ describe("resolveApiKeyForProfile fallback to main agent", () => {
   afterEach(async () => {
     vi.unstubAllGlobals();
 
-    // Restore original environment
-    if (previousStateDir === undefined) {
-      delete process.env.OPENCLAW_STATE_DIR;
-    } else {
-      process.env.OPENCLAW_STATE_DIR = previousStateDir;
-    }
-    if (previousAgentDir === undefined) {
-      delete process.env.OPENCLAW_AGENT_DIR;
-    } else {
-      process.env.OPENCLAW_AGENT_DIR = previousAgentDir;
-    }
-    if (previousPiAgentDir === undefined) {
-      delete process.env.PI_CODING_AGENT_DIR;
-    } else {
-      process.env.PI_CODING_AGENT_DIR = previousPiAgentDir;
-    }
+    envSnapshot.restore();
 
     await fs.rm(tmpDir, { recursive: true, force: true });
   });

src/agents/pi-tools.safe-bins.e2e.test.ts

@@ -4,8 +4,9 @@ import path from "node:path";
 import { afterAll, beforeAll, describe, expect, it, vi } from "vitest";
 import type { OpenClawConfig } from "../config/config.js";
 import type { ExecApprovalsResolved } from "../infra/exec-approvals.js";
+import { captureEnv } from "../test-utils/env.js";
 
-const previousBundledPluginsDir = process.env.OPENCLAW_BUNDLED_PLUGINS_DIR;
+const bundledPluginsDirSnapshot = captureEnv(["OPENCLAW_BUNDLED_PLUGINS_DIR"]);
 
 beforeAll(() => {
   process.env.OPENCLAW_BUNDLED_PLUGINS_DIR = path.join(
@@ -15,11 +16,7 @@ beforeAll(() => {
 });
 
 afterAll(() => {
-  if (previousBundledPluginsDir === undefined) {
-    delete process.env.OPENCLAW_BUNDLED_PLUGINS_DIR;
-  } else {
-    process.env.OPENCLAW_BUNDLED_PLUGINS_DIR = previousBundledPluginsDir;
-  }
+  bundledPluginsDirSnapshot.restore();
 });
 
 vi.mock("../infra/shell-env.js", async (importOriginal) => {
@@ -109,20 +106,16 @@ describe("createOpenClawCodingTools safeBins", () => {
     expect(execTool).toBeDefined();
 
     const marker = `safe-bins-${Date.now()}`;
-    const prevShellEnvTimeoutMs = process.env.OPENCLAW_SHELL_ENV_TIMEOUT_MS;
-    process.env.OPENCLAW_SHELL_ENV_TIMEOUT_MS = "1000";
+    const envSnapshot = captureEnv(["OPENCLAW_SHELL_ENV_TIMEOUT_MS"]);
     const result = await (async () => {
       try {
+        process.env.OPENCLAW_SHELL_ENV_TIMEOUT_MS = "1000";
         return await execTool!.execute("call1", {
           command: `echo ${marker}`,
           workdir: tmpDir,
         });
       } finally {
-        if (prevShellEnvTimeoutMs === undefined) {
-          delete process.env.OPENCLAW_SHELL_ENV_TIMEOUT_MS;
-        } else {
-          process.env.OPENCLAW_SHELL_ENV_TIMEOUT_MS = prevShellEnvTimeoutMs;
-        }
+        envSnapshot.restore();
       }
     })();
     const text = result.content.find((content) => content.type === "text")?.text ?? "";

src/test-helpers/state-dir-env.ts

@@ -1,26 +1,11 @@
-type StateDirEnvSnapshot = {
-  openclawStateDir: string | undefined;
-  clawdbotStateDir: string | undefined;
-};
+import { captureEnv } from "../test-utils/env.js";
 
-export function snapshotStateDirEnv(): StateDirEnvSnapshot {
-  return {
-    openclawStateDir: process.env.OPENCLAW_STATE_DIR,
-    clawdbotStateDir: process.env.CLAWDBOT_STATE_DIR,
-  };
+export function snapshotStateDirEnv() {
+  return captureEnv(["OPENCLAW_STATE_DIR", "CLAWDBOT_STATE_DIR"]);
 }
 
-export function restoreStateDirEnv(snapshot: StateDirEnvSnapshot): void {
-  if (snapshot.openclawStateDir === undefined) {
-    delete process.env.OPENCLAW_STATE_DIR;
-  } else {
-    process.env.OPENCLAW_STATE_DIR = snapshot.openclawStateDir;
-  }
-  if (snapshot.clawdbotStateDir === undefined) {
-    delete process.env.CLAWDBOT_STATE_DIR;
-  } else {
-    process.env.CLAWDBOT_STATE_DIR = snapshot.clawdbotStateDir;
-  }
+export function restoreStateDirEnv(snapshot: ReturnType<typeof snapshotStateDirEnv>): void {
+  snapshot.restore();
 }
 
 export function setStateDirEnv(stateDir: string): void {