From f0e373b82ee344ae0bb6b1f4b70adb8a141db86a Mon Sep 17 00:00:00 2001
From: Peter Steinberger <steipete@gmail.com>
Date: Sun, 15 Feb 2026 23:29:56 +0000
Subject: [PATCH] refactor(test): simplify state dir env restore

---
 src/media/store.redirect.test.ts | 11 +++++------
 src/web/media.test.ts            | 11 ++++-------
 2 files changed, 9 insertions(+), 13 deletions(-)

diff --git a/src/media/store.redirect.test.ts b/src/media/store.redirect.test.ts
index 40ba39815..3d87f5bf9 100644
--- a/src/media/store.redirect.test.ts
+++ b/src/media/store.redirect.test.ts
@@ -4,14 +4,17 @@ import os from "node:os";
 import path from "node:path";
 import { PassThrough } from "node:stream";
 import { afterAll, beforeAll, beforeEach, describe, expect, it, vi } from "vitest";
+import { captureEnv } from "../test-utils/env.js";
 import { saveMediaSource, setMediaStoreNetworkDepsForTest } from "./store.js";
 
 const HOME = path.join(os.tmpdir(), "openclaw-home-redirect");
-const previousStateDir = process.env.OPENCLAW_STATE_DIR;
 const mockRequest = vi.fn();
 
 describe("media store redirects", () => {
+  let envSnapshot: ReturnType<typeof captureEnv>;
+
   beforeAll(async () => {
+    envSnapshot = captureEnv(["OPENCLAW_STATE_DIR"]);
     await fs.rm(HOME, { recursive: true, force: true });
     process.env.OPENCLAW_STATE_DIR = HOME;
   });
@@ -29,11 +32,7 @@ describe("media store redirects", () => {
 
   afterAll(async () => {
     await fs.rm(HOME, { recursive: true, force: true });
-    if (previousStateDir === undefined) {
-      delete process.env.OPENCLAW_STATE_DIR;
-    } else {
-      process.env.OPENCLAW_STATE_DIR = previousStateDir;
-    }
+    envSnapshot.restore();
     setMediaStoreNetworkDepsForTest();
     vi.clearAllMocks();
   });
diff --git a/src/web/media.test.ts b/src/web/media.test.ts
index c76fb55d8..fadcea5cc 100644
--- a/src/web/media.test.ts
+++ b/src/web/media.test.ts
@@ -5,6 +5,7 @@ import sharp from "sharp";
 import { afterAll, afterEach, beforeAll, describe, expect, it, vi } from "vitest";
 import * as ssrf from "../infra/net/ssrf.js";
 import { optimizeImageToPng } from "../media/image-ops.js";
+import { captureEnv } from "../test-utils/env.js";
 import { loadWebMedia, loadWebMediaRaw, optimizeImageToJpeg } from "./media.js";
 
 let fixtureRoot = "";
@@ -19,7 +20,7 @@ let alphaPngFile = "";
 let fallbackPngBuffer: Buffer;
 let fallbackPngFile = "";
 let fallbackPngCap = 0;
-let previousStateDir: string | undefined;
+let stateDirSnapshot: ReturnType<typeof captureEnv>;
 
 async function writeTempFile(buffer: Buffer, ext: string): Promise<string> {
   const file = path.join(fixtureRoot, `media-${fixtureFileCount++}${ext}`);
@@ -100,7 +101,7 @@ describe("web media loading", () => {
   beforeAll(() => {
     // Ensure state dir is stable and not influenced by other tests that stub OPENCLAW_STATE_DIR.
     // Also keep it outside os.tmpdir() so tmpdir localRoots doesn't accidentally make all state readable.
-    previousStateDir = process.env.OPENCLAW_STATE_DIR;
+    stateDirSnapshot = captureEnv(["OPENCLAW_STATE_DIR"]);
     process.env.OPENCLAW_STATE_DIR = path.join(
       path.parse(os.tmpdir()).root,
       "var",
@@ -110,11 +111,7 @@ describe("web media loading", () => {
   });
 
   afterAll(() => {
-    if (previousStateDir === undefined) {
-      delete process.env.OPENCLAW_STATE_DIR;
-    } else {
-      process.env.OPENCLAW_STATE_DIR = previousStateDir;
-    }
+    stateDirSnapshot.restore();
   });
 
   beforeAll(() => {