From 7675b10223d1548cb2b268c35c5a2dcb60c0add7 Mon Sep 17 00:00:00 2001
From: Shakker <shakkerdroid@gmail.com>
Date: Fri, 5 Jun 2026 00:15:44 +0100
Subject: [PATCH] test: capture usage format env setup

---
 src/utils/usage-format.test.ts | 17 +++++------------
 1 file changed, 5 insertions(+), 12 deletions(-)

diff --git a/src/utils/usage-format.test.ts b/src/utils/usage-format.test.ts
index 6ca9caeba621..f237329b0baa 100644
--- a/src/utils/usage-format.test.ts
+++ b/src/utils/usage-format.test.ts
@@ -9,6 +9,7 @@ import {
   setGatewayModelPricingForTest,
 } from "../gateway/model-pricing-cache-state.js";
 import * as manifestModelIdNormalization from "../plugins/manifest-model-id-normalization.js";
+import { captureEnv } from "../test-utils/env.js";
 import {
   resetUsageFormatCachesForTest,
   estimateUsageCost,
@@ -42,12 +43,12 @@ function requireTieredPricing(
 }
 
 describe("usage-format", () => {
-  const originalAgentDir = process.env.OPENCLAW_AGENT_DIR;
-  const originalStateDir = process.env.OPENCLAW_STATE_DIR;
+  let envSnapshot: ReturnType<typeof captureEnv> | undefined;
   let agentDir: string;
   let stateDir: string;
 
   beforeEach(async () => {
+    envSnapshot = captureEnv(["OPENCLAW_AGENT_DIR", "OPENCLAW_STATE_DIR"]);
     stateDir = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-usage-format-"));
     agentDir = path.join(stateDir, "agents", "main", "agent");
     process.env.OPENCLAW_STATE_DIR = stateDir;
@@ -58,16 +59,8 @@ describe("usage-format", () => {
   });
 
   afterEach(async () => {
-    if (originalAgentDir === undefined) {
-      delete process.env.OPENCLAW_AGENT_DIR;
-    } else {
-      process.env.OPENCLAW_AGENT_DIR = originalAgentDir;
-    }
-    if (originalStateDir === undefined) {
-      delete process.env.OPENCLAW_STATE_DIR;
-    } else {
-      process.env.OPENCLAW_STATE_DIR = originalStateDir;
-    }
+    envSnapshot?.restore();
+    envSnapshot = undefined;
     resetUsageFormatCachesForTest();
     resetGatewayModelPricingCacheForTest();
     await fs.rm(stateDir, { recursive: true, force: true });