From cf8261425954692c75346c32c1f5d115cdab5df1 Mon Sep 17 00:00:00 2001
From: Peter Steinberger <steipete@gmail.com>
Date: Sat, 21 Feb 2026 13:02:12 +0000
Subject: [PATCH] refactor(test): reuse state-dir helper in telegram tests

---
 src/telegram/token.test.ts               | 20 ++----------------
 src/telegram/update-offset-store.test.ts | 26 ++++--------------------
 src/test-helpers/state-dir-env.test.ts   | 20 ++++++++++++++++++
 3 files changed, 26 insertions(+), 40 deletions(-)

diff --git a/src/telegram/token.test.ts b/src/telegram/token.test.ts
index 34b06ca1d44..69a9e8aa1b8 100644
--- a/src/telegram/token.test.ts
+++ b/src/telegram/token.test.ts
@@ -1,9 +1,9 @@
 import fs from "node:fs";
-import fsPromises from "node:fs/promises";
 import os from "node:os";
 import path from "node:path";
 import { afterEach, describe, expect, it, vi } from "vitest";
 import type { OpenClawConfig } from "../config/config.js";
+import { withStateDirEnv } from "../test-helpers/state-dir-env.js";
 import { resolveTelegramToken } from "./token.js";
 import { readTelegramUpdateOffset, writeTelegramUpdateOffset } from "./update-offset-store.js";
 
@@ -11,22 +11,6 @@ function withTempDir(): string {
   return fs.mkdtempSync(path.join(os.tmpdir(), "openclaw-telegram-token-"));
 }
 
-async function withTempStateDir<T>(fn: (dir: string) => Promise<T>) {
-  const previous = process.env.OPENCLAW_STATE_DIR;
-  const dir = await fsPromises.mkdtemp(path.join(os.tmpdir(), "openclaw-telegram-"));
-  process.env.OPENCLAW_STATE_DIR = dir;
-  try {
-    return await fn(dir);
-  } finally {
-    if (previous === undefined) {
-      delete process.env.OPENCLAW_STATE_DIR;
-    } else {
-      process.env.OPENCLAW_STATE_DIR = previous;
-    }
-    await fsPromises.rm(dir, { recursive: true, force: true });
-  }
-}
-
 describe("resolveTelegramToken", () => {
   afterEach(() => {
     vi.unstubAllEnvs();
@@ -108,7 +92,7 @@ describe("resolveTelegramToken", () => {
 
 describe("telegram update offset store", () => {
   it("persists and reloads the last update id", async () => {
-    await withTempStateDir(async () => {
+    await withStateDirEnv("openclaw-telegram-", async () => {
       expect(await readTelegramUpdateOffset({ accountId: "primary" })).toBeNull();
 
       await writeTelegramUpdateOffset({
diff --git a/src/telegram/update-offset-store.test.ts b/src/telegram/update-offset-store.test.ts
index 38fd9753d7c..523038b30f8 100644
--- a/src/telegram/update-offset-store.test.ts
+++ b/src/telegram/update-offset-store.test.ts
@@ -1,32 +1,14 @@
-import fs from "node:fs/promises";
-import os from "node:os";
-import path from "node:path";
 import { describe, expect, it } from "vitest";
+import { withStateDirEnv } from "../test-helpers/state-dir-env.js";
 import {
   deleteTelegramUpdateOffset,
   readTelegramUpdateOffset,
   writeTelegramUpdateOffset,
 } from "./update-offset-store.js";
 
-async function withTempStateDir<T>(fn: (dir: string) => Promise<T>) {
-  const previous = process.env.OPENCLAW_STATE_DIR;
-  const dir = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-tg-offset-"));
-  process.env.OPENCLAW_STATE_DIR = dir;
-  try {
-    return await fn(dir);
-  } finally {
-    if (previous === undefined) {
-      delete process.env.OPENCLAW_STATE_DIR;
-    } else {
-      process.env.OPENCLAW_STATE_DIR = previous;
-    }
-    await fs.rm(dir, { recursive: true, force: true });
-  }
-}
-
 describe("deleteTelegramUpdateOffset", () => {
   it("removes the offset file so a new bot starts fresh", async () => {
-    await withTempStateDir(async () => {
+    await withStateDirEnv("openclaw-tg-offset-", async () => {
       await writeTelegramUpdateOffset({ accountId: "default", updateId: 432_000_000 });
       expect(await readTelegramUpdateOffset({ accountId: "default" })).toBe(432_000_000);
 
@@ -36,13 +18,13 @@ describe("deleteTelegramUpdateOffset", () => {
   });
 
   it("does not throw when the offset file does not exist", async () => {
-    await withTempStateDir(async () => {
+    await withStateDirEnv("openclaw-tg-offset-", async () => {
       await expect(deleteTelegramUpdateOffset({ accountId: "nonexistent" })).resolves.not.toThrow();
     });
   });
 
   it("only removes the targeted account offset, leaving others intact", async () => {
-    await withTempStateDir(async () => {
+    await withStateDirEnv("openclaw-tg-offset-", async () => {
       await writeTelegramUpdateOffset({ accountId: "default", updateId: 100 });
       await writeTelegramUpdateOffset({ accountId: "alerts", updateId: 200 });
 
diff --git a/src/test-helpers/state-dir-env.test.ts b/src/test-helpers/state-dir-env.test.ts
index 744e42f868f..e251609c8c0 100644
--- a/src/test-helpers/state-dir-env.test.ts
+++ b/src/test-helpers/state-dir-env.test.ts
@@ -42,4 +42,24 @@ describe("state-dir-env helpers", () => {
     await expect(fs.stat(capturedStateDir)).rejects.toThrow();
     await expect(fs.stat(capturedTempRoot)).rejects.toThrow();
   });
+
+  it("withStateDirEnv restores env and cleans temp root when callback throws", async () => {
+    const prevOpenClaw = process.env.OPENCLAW_STATE_DIR;
+    const prevLegacy = process.env.CLAWDBOT_STATE_DIR;
+
+    let capturedTempRoot = "";
+    let capturedStateDir = "";
+    await expect(
+      withStateDirEnv("openclaw-state-dir-env-", async ({ tempRoot, stateDir }) => {
+        capturedTempRoot = tempRoot;
+        capturedStateDir = stateDir;
+        throw new Error("boom");
+      }),
+    ).rejects.toThrow("boom");
+
+    expect(process.env.OPENCLAW_STATE_DIR).toBe(prevOpenClaw);
+    expect(process.env.CLAWDBOT_STATE_DIR).toBe(prevLegacy);
+    await expect(fs.stat(capturedStateDir)).rejects.toThrow();
+    await expect(fs.stat(capturedTempRoot)).rejects.toThrow();
+  });
 });