test: dedupe telegram draft stream setup and extend state-dir env coverage

src/telegram/draft-stream.test.ts

@@ -2,6 +2,8 @@ import type { Bot } from "grammy";
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 import { createTelegramDraftStream } from "./draft-stream.js";
 
+type TelegramDraftStreamParams = Parameters<typeof createTelegramDraftStream>[0];
+
 function createMockDraftApi(sendMessageImpl?: () => Promise<{ message_id: number }>) {
   return {
     sendMessage: vi.fn(sendMessageImpl ?? (async () => ({ message_id: 17 }))),
@@ -17,11 +19,18 @@ function createForumDraftStream(api: ReturnType<typeof createMockDraftApi>) {
 function createThreadedDraftStream(
   api: ReturnType<typeof createMockDraftApi>,
   thread: { id: number; scope: "forum" | "dm" },
+) {
+  return createDraftStream(api, { thread });
+}
+
+function createDraftStream(
+  api: ReturnType<typeof createMockDraftApi>,
+  overrides: Omit<Partial<TelegramDraftStreamParams>, "api" | "chatId"> = {},
 ) {
   return createTelegramDraftStream({
     api: api as unknown as Bot["api"],
     chatId: 123,
-    thread,
+    ...overrides,
   });
 }
 
@@ -34,6 +43,18 @@ async function expectInitialForumSend(
   );
 }
 
+function createForceNewMessageHarness(params: { throttleMs?: number } = {}) {
+  const api = createMockDraftApi();
+  api.sendMessage
+    .mockResolvedValueOnce({ message_id: 17 })
+    .mockResolvedValueOnce({ message_id: 42 });
+  const stream = createDraftStream(
+    api,
+    params.throttleMs != null ? { throttleMs: params.throttleMs } : {},
+  );
+  return { api, stream };
+}
+
 describe("createTelegramDraftStream", () => {
   it("sends stream preview message with message_thread_id when provided", async () => {
     const api = createMockDraftApi();
@@ -100,18 +121,7 @@ describe("createTelegramDraftStream", () => {
   });
 
   it("creates new message after forceNewMessage is called", async () => {
-    const api = {
-      sendMessage: vi
-        .fn()
-        .mockResolvedValueOnce({ message_id: 17 })
-        .mockResolvedValueOnce({ message_id: 42 }),
-      editMessageText: vi.fn().mockResolvedValue(true),
-      deleteMessage: vi.fn().mockResolvedValue(true),
-    };
-    const stream = createTelegramDraftStream({
-      api: api as unknown as Bot["api"],
-      chatId: 123,
-    });
+    const { api, stream } = createForceNewMessageHarness();
 
     // First message
     stream.update("Hello");
@@ -136,19 +146,7 @@ describe("createTelegramDraftStream", () => {
   it("sends first update immediately after forceNewMessage within throttle window", async () => {
     vi.useFakeTimers();
     try {
-      const api = {
-        sendMessage: vi
-          .fn()
-          .mockResolvedValueOnce({ message_id: 17 })
-          .mockResolvedValueOnce({ message_id: 42 }),
-        editMessageText: vi.fn().mockResolvedValue(true),
-        deleteMessage: vi.fn().mockResolvedValue(true),
-      };
-      const stream = createTelegramDraftStream({
-        api: api as unknown as Bot["api"],
-        chatId: 123,
-        throttleMs: 1000,
-      });
+      const { api, stream } = createForceNewMessageHarness({ throttleMs: 1000 });
 
       stream.update("Hello");
       await vi.waitFor(() => expect(api.sendMessage).toHaveBeenCalledTimes(1));

src/test-helpers/state-dir-env.test.ts

@@ -29,6 +29,16 @@ async function expectPathMissing(filePath: string) {
   await expect(fs.stat(filePath)).rejects.toThrow();
 }
 
+async function expectStateDirEnvRestored(params: {
+  prev: EnvSnapshot;
+  capturedStateDir: string;
+  capturedTempRoot: string;
+}) {
+  expectStateDirVars(params.prev);
+  await expectPathMissing(params.capturedStateDir);
+  await expectPathMissing(params.capturedTempRoot);
+}
+
 describe("state-dir-env helpers", () => {
   it("set/snapshot/restore round-trips OPENCLAW_STATE_DIR", () => {
     const prev = snapshotCurrentStateDirVars();
@@ -55,9 +65,7 @@ describe("state-dir-env helpers", () => {
       await fs.writeFile(path.join(stateDir, "probe.txt"), "ok", "utf8");
     });
 
-    expectStateDirVars(prev);
-    await expectPathMissing(capturedStateDir);
-    await expectPathMissing(capturedTempRoot);
+    await expectStateDirEnvRestored({ prev, capturedStateDir, capturedTempRoot });
   });
 
   it("withStateDirEnv restores env and cleans temp root when callback throws", async () => {
@@ -73,8 +81,28 @@ describe("state-dir-env helpers", () => {
       }),
     ).rejects.toThrow("boom");
 
-    expectStateDirVars(prev);
-    await expectPathMissing(capturedStateDir);
-    await expectPathMissing(capturedTempRoot);
+    await expectStateDirEnvRestored({ prev, capturedStateDir, capturedTempRoot });
+  });
+
+  it("withStateDirEnv restores both env vars when legacy var was previously set", async () => {
+    const testSnapshot = snapshotStateDirEnv();
+    process.env.OPENCLAW_STATE_DIR = "/tmp/original-openclaw";
+    process.env.CLAWDBOT_STATE_DIR = "/tmp/original-legacy";
+    const prev = snapshotCurrentStateDirVars();
+
+    let capturedTempRoot = "";
+    let capturedStateDir = "";
+    try {
+      await withStateDirEnv("openclaw-state-dir-env-", async ({ tempRoot, stateDir }) => {
+        capturedTempRoot = tempRoot;
+        capturedStateDir = stateDir;
+        expect(process.env.OPENCLAW_STATE_DIR).toBe(stateDir);
+        expect(process.env.CLAWDBOT_STATE_DIR).toBeUndefined();
+      });
+
+      await expectStateDirEnvRestored({ prev, capturedStateDir, capturedTempRoot });
+    } finally {
+      restoreStateDirEnv(testSnapshot);
+    }
   });
 });