From f809ff5e55a63a40b1928962900b1bc352db27ef Mon Sep 17 00:00:00 2001
From: Peter Steinberger <steipete@gmail.com>
Date: Sun, 15 Feb 2026 23:51:24 +0000
Subject: [PATCH] refactor(test): reuse env snapshot helper

---
 src/agents/auth-profiles.chutes.e2e.test.ts   | 34 ++++++-------------
 .../subagent-registry.persistence.e2e.test.ts |  9 ++---
 .../usage.sessions-usage.test.ts              |  9 ++---
 3 files changed, 16 insertions(+), 36 deletions(-)

diff --git a/src/agents/auth-profiles.chutes.e2e.test.ts b/src/agents/auth-profiles.chutes.e2e.test.ts
index 317ce9c771..c21f37ed1c 100644
--- a/src/agents/auth-profiles.chutes.e2e.test.ts
+++ b/src/agents/auth-profiles.chutes.e2e.test.ts
@@ -2,6 +2,7 @@ import fs from "node:fs/promises";
 import os from "node:os";
 import path from "node:path";
 import { afterEach, describe, expect, it, vi } from "vitest";
+import { captureEnv } from "../test-utils/env.js";
 import {
   type AuthProfileStore,
   ensureAuthProfileStore,
@@ -10,10 +11,7 @@ import {
 import { CHUTES_TOKEN_ENDPOINT, type ChutesStoredOAuth } from "./chutes-oauth.js";
 
 describe("auth-profiles (chutes)", () => {
-  const previousStateDir = process.env.OPENCLAW_STATE_DIR;
-  const previousAgentDir = process.env.OPENCLAW_AGENT_DIR;
-  const previousPiAgentDir = process.env.PI_CODING_AGENT_DIR;
-  const previousChutesClientId = process.env.CHUTES_CLIENT_ID;
+  let envSnapshot: ReturnType<typeof captureEnv> | undefined;
   let tempDir: string | null = null;
 
   afterEach(async () => {
@@ -22,29 +20,17 @@ describe("auth-profiles (chutes)", () => {
       await fs.rm(tempDir, { recursive: true, force: true });
       tempDir = null;
     }
-    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;
-    }
-    if (previousChutesClientId === undefined) {
-      delete process.env.CHUTES_CLIENT_ID;
-    } else {
-      process.env.CHUTES_CLIENT_ID = previousChutesClientId;
-    }
+    envSnapshot?.restore();
+    envSnapshot = undefined;
   });
 
   it("refreshes expired Chutes OAuth credentials", async () => {
+    envSnapshot = captureEnv([
+      "OPENCLAW_STATE_DIR",
+      "OPENCLAW_AGENT_DIR",
+      "PI_CODING_AGENT_DIR",
+      "CHUTES_CLIENT_ID",
+    ]);
     tempDir = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-chutes-"));
     process.env.OPENCLAW_STATE_DIR = tempDir;
     process.env.OPENCLAW_AGENT_DIR = path.join(tempDir, "agents", "main", "agent");
diff --git a/src/agents/subagent-registry.persistence.e2e.test.ts b/src/agents/subagent-registry.persistence.e2e.test.ts
index 9b3f5348c4..4a6620c4e5 100644
--- a/src/agents/subagent-registry.persistence.e2e.test.ts
+++ b/src/agents/subagent-registry.persistence.e2e.test.ts
@@ -2,6 +2,7 @@ import fs from "node:fs/promises";
 import os from "node:os";
 import path from "node:path";
 import { afterEach, describe, expect, it, vi } from "vitest";
+import { captureEnv } from "../test-utils/env.js";
 import {
   initSubagentRegistry,
   registerSubagentRun,
@@ -29,7 +30,7 @@ vi.mock("./subagent-announce.js", () => ({
 }));
 
 describe("subagent registry persistence", () => {
-  const previousStateDir = process.env.OPENCLAW_STATE_DIR;
+  const envSnapshot = captureEnv(["OPENCLAW_STATE_DIR"]);
   let tempStateDir: string | null = null;
 
   afterEach(async () => {
@@ -39,11 +40,7 @@ describe("subagent registry persistence", () => {
       await fs.rm(tempStateDir, { recursive: true, force: true });
       tempStateDir = null;
     }
-    if (previousStateDir === undefined) {
-      delete process.env.OPENCLAW_STATE_DIR;
-    } else {
-      process.env.OPENCLAW_STATE_DIR = previousStateDir;
-    }
+    envSnapshot.restore();
   });
 
   it("persists runs to disk and resumes after restart", async () => {
diff --git a/src/gateway/server-methods/usage.sessions-usage.test.ts b/src/gateway/server-methods/usage.sessions-usage.test.ts
index efdeb9a164..3731bdd4e0 100644
--- a/src/gateway/server-methods/usage.sessions-usage.test.ts
+++ b/src/gateway/server-methods/usage.sessions-usage.test.ts
@@ -2,6 +2,7 @@ import fs from "node:fs";
 import os from "node:os";
 import path from "node:path";
 import { beforeEach, describe, expect, it, vi } from "vitest";
+import { captureEnv } from "../../test-utils/env.js";
 
 vi.mock("../../config/config.js", () => {
   return {
@@ -118,7 +119,7 @@ describe("sessions.usage", () => {
   it("resolves store entries by sessionId when queried via discovered agent-prefixed key", async () => {
     const storeKey = "agent:opus:slack:dm:u123";
     const stateDir = fs.mkdtempSync(path.join(os.tmpdir(), "openclaw-usage-test-"));
-    const previousStateDir = process.env.OPENCLAW_STATE_DIR;
+    const envSnapshot = captureEnv(["OPENCLAW_STATE_DIR"]);
     process.env.OPENCLAW_STATE_DIR = stateDir;
 
     try {
@@ -163,11 +164,7 @@ describe("sessions.usage", () => {
         vi.mocked(loadSessionCostSummary).mock.calls.some((call) => call[0]?.agentId === "opus"),
       ).toBe(true);
     } finally {
-      if (previousStateDir === undefined) {
-        delete process.env.OPENCLAW_STATE_DIR;
-      } else {
-        process.env.OPENCLAW_STATE_DIR = previousStateDir;
-      }
+      envSnapshot.restore();
       fs.rmSync(stateDir, { recursive: true, force: true });
     }
   });