From bd9d3e2f87d32bb1baca4444fb959807fa2f711c Mon Sep 17 00:00:00 2001
From: Peter Steinberger <steipete@gmail.com>
Date: Sat, 21 Feb 2026 18:24:24 +0000
Subject: [PATCH] refactor(test): reuse env helper in update cli tests

---
 src/cli/update-cli.test.ts | 45 ++++++++++++++++++--------------------
 1 file changed, 21 insertions(+), 24 deletions(-)

diff --git a/src/cli/update-cli.test.ts b/src/cli/update-cli.test.ts
index 85a3dac2d..2a5cb8f48 100644
--- a/src/cli/update-cli.test.ts
+++ b/src/cli/update-cli.test.ts
@@ -3,7 +3,7 @@ import path from "node:path";
 import { beforeEach, describe, expect, it, vi } from "vitest";
 import type { OpenClawConfig, ConfigFileSnapshot } from "../config/types.openclaw.js";
 import type { UpdateRunResult } from "../infra/update-runner.js";
-import { captureEnv } from "../test-utils/env.js";
+import { withEnvAsync } from "../test-utils/env.js";
 
 const confirm = vi.fn();
 const select = vi.fn();
@@ -604,30 +604,31 @@ describe("update-cli", () => {
   });
 
   it("updateCommand continues after doctor sub-step and clears update flag", async () => {
-    const envSnapshot = captureEnv(["OPENCLAW_UPDATE_IN_PROGRESS"]);
     const randomSpy = vi.spyOn(Math, "random").mockReturnValue(0);
     try {
-      delete process.env.OPENCLAW_UPDATE_IN_PROGRESS;
-      vi.mocked(runGatewayUpdate).mockResolvedValue(makeOkUpdateResult());
-      vi.mocked(runDaemonRestart).mockResolvedValue(true);
-      vi.mocked(doctorCommand).mockResolvedValue(undefined);
-      vi.mocked(defaultRuntime.log).mockClear();
+      await withEnvAsync({ OPENCLAW_UPDATE_IN_PROGRESS: undefined }, async () => {
+        vi.mocked(runGatewayUpdate).mockResolvedValue(makeOkUpdateResult());
+        vi.mocked(runDaemonRestart).mockResolvedValue(true);
+        vi.mocked(doctorCommand).mockResolvedValue(undefined);
+        vi.mocked(defaultRuntime.log).mockClear();
 
-      await updateCommand({});
+        await updateCommand({});
 
-      expect(doctorCommand).toHaveBeenCalledWith(
-        defaultRuntime,
-        expect.objectContaining({ nonInteractive: true }),
-      );
-      expect(process.env.OPENCLAW_UPDATE_IN_PROGRESS).toBeUndefined();
+        expect(doctorCommand).toHaveBeenCalledWith(
+          defaultRuntime,
+          expect.objectContaining({ nonInteractive: true }),
+        );
+        expect(process.env.OPENCLAW_UPDATE_IN_PROGRESS).toBeUndefined();
 
-      const logLines = vi.mocked(defaultRuntime.log).mock.calls.map((call) => String(call[0]));
-      expect(
-        logLines.some((line) => line.includes("Leveled up! New skills unlocked. You're welcome.")),
-      ).toBe(true);
+        const logLines = vi.mocked(defaultRuntime.log).mock.calls.map((call) => String(call[0]));
+        expect(
+          logLines.some((line) =>
+            line.includes("Leveled up! New skills unlocked. You're welcome."),
+          ),
+        ).toBe(true);
+      });
     } finally {
       randomSpy.mockRestore();
-      envSnapshot.restore();
     }
   });
 
@@ -731,10 +732,8 @@ describe("update-cli", () => {
 
   it("updateWizardCommand offers dev checkout and forwards selections", async () => {
     const tempDir = createCaseDir("openclaw-update-wizard");
-    const envSnapshot = captureEnv(["OPENCLAW_GIT_DIR"]);
-    try {
+    await withEnvAsync({ OPENCLAW_GIT_DIR: tempDir }, async () => {
       setTty(true);
-      process.env.OPENCLAW_GIT_DIR = tempDir;
 
       vi.mocked(checkUpdateStatus).mockResolvedValue({
         root: "/test/path",
@@ -760,8 +759,6 @@ describe("update-cli", () => {
 
       const call = vi.mocked(runGatewayUpdate).mock.calls[0]?.[0];
       expect(call?.channel).toBe("dev");
-    } finally {
-      envSnapshot.restore();
-    }
+    });
   });
 });