refactor(test): dedupe telegram draft-stream fixtures

src/telegram/bot-message-dispatch.test.ts

@@ -2,6 +2,10 @@ import path from "node:path";
 import type { Bot } from "grammy";
 import { beforeEach, describe, expect, it, vi } from "vitest";
 import { STATE_DIR } from "../config/paths.js";
+import {
+  createSequencedTestDraftStream,
+  createTestDraftStream,
+} from "./draft-stream.test-helpers.js";
 
 const createTelegramDraftStream = vi.hoisted(() => vi.fn());
 const dispatchReplyWithBufferedBlockDispatcher = vi.hoisted(() => vi.fn());
@@ -52,44 +56,9 @@ describe("dispatchTelegramMessage draft streaming", () => {
     loadSessionStore.mockReturnValue({});
   });
 
-  function createDraftStream(messageId?: number) {
+  const createDraftStream = (messageId?: number) => createTestDraftStream({ messageId });
-    let previewRevision = 0;
+  const createSequencedDraftStream = (startMessageId = 1001) =>
-    return {
+    createSequencedTestDraftStream(startMessageId);
-      update: vi.fn().mockImplementation(() => {
-        previewRevision += 1;
-      }),
-      flush: vi.fn().mockResolvedValue(true),
-      messageId: vi.fn().mockReturnValue(messageId),
-      previewMode: vi.fn().mockReturnValue("message"),
-      previewRevision: vi.fn().mockImplementation(() => previewRevision),
-      clear: vi.fn().mockResolvedValue(undefined),
-      stop: vi.fn().mockResolvedValue(undefined),
-      forceNewMessage: vi.fn(),
-    };
-  }
-
-  function createSequencedDraftStream(startMessageId = 1001) {
-    let activeMessageId: number | undefined;
-    let nextMessageId = startMessageId;
-    let previewRevision = 0;
-    return {
-      update: vi.fn().mockImplementation(() => {
-        if (activeMessageId == null) {
-          activeMessageId = nextMessageId++;
-        }
-        previewRevision += 1;
-      }),
-      flush: vi.fn().mockResolvedValue(true),
-      messageId: vi.fn().mockImplementation(() => activeMessageId),
-      previewMode: vi.fn().mockReturnValue("message"),
-      previewRevision: vi.fn().mockImplementation(() => previewRevision),
-      clear: vi.fn().mockResolvedValue(undefined),
-      stop: vi.fn().mockResolvedValue(undefined),
-      forceNewMessage: vi.fn().mockImplementation(() => {
-        activeMessageId = undefined;
-      }),
-    };
-  }
 
   function setupDraftStreams(params?: { answerMessageId?: number; reasoningMessageId?: number }) {
     const answerDraftStream = createDraftStream(params?.answerMessageId);

src/telegram/draft-stream.test-helpers.ts

@@ -0,0 +1,74 @@
+import { vi } from "vitest";
+
+type DraftPreviewMode = "message" | "draft";
+
+export type TestDraftStream = {
+  update: ReturnType<typeof vi.fn<(text: string) => void>>;
+  flush: ReturnType<typeof vi.fn<() => Promise<void>>>;
+  messageId: ReturnType<typeof vi.fn<() => number | undefined>>;
+  previewMode: ReturnType<typeof vi.fn<() => DraftPreviewMode>>;
+  previewRevision: ReturnType<typeof vi.fn<() => number>>;
+  clear: ReturnType<typeof vi.fn<() => Promise<void>>>;
+  stop: ReturnType<typeof vi.fn<() => Promise<void>>>;
+  forceNewMessage: ReturnType<typeof vi.fn<() => void>>;
+  setMessageId: (value: number | undefined) => void;
+};
+
+export function createTestDraftStream(params?: {
+  messageId?: number;
+  previewMode?: DraftPreviewMode;
+  onUpdate?: (text: string) => void;
+  onStop?: () => void | Promise<void>;
+  clearMessageIdOnForceNew?: boolean;
+}): TestDraftStream {
+  let messageId = params?.messageId;
+  let previewRevision = 0;
+  return {
+    update: vi.fn().mockImplementation((text: string) => {
+      previewRevision += 1;
+      params?.onUpdate?.(text);
+    }),
+    flush: vi.fn().mockResolvedValue(undefined),
+    messageId: vi.fn().mockImplementation(() => messageId),
+    previewMode: vi.fn().mockReturnValue(params?.previewMode ?? "message"),
+    previewRevision: vi.fn().mockImplementation(() => previewRevision),
+    clear: vi.fn().mockResolvedValue(undefined),
+    stop: vi.fn().mockImplementation(async () => {
+      await params?.onStop?.();
+    }),
+    forceNewMessage: vi.fn().mockImplementation(() => {
+      if (params?.clearMessageIdOnForceNew) {
+        messageId = undefined;
+      }
+    }),
+    setMessageId: (value: number | undefined) => {
+      messageId = value;
+    },
+  };
+}
+
+export function createSequencedTestDraftStream(startMessageId = 1001): TestDraftStream {
+  let activeMessageId: number | undefined;
+  let nextMessageId = startMessageId;
+  let previewRevision = 0;
+  return {
+    update: vi.fn().mockImplementation(() => {
+      if (activeMessageId == null) {
+        activeMessageId = nextMessageId++;
+      }
+      previewRevision += 1;
+    }),
+    flush: vi.fn().mockResolvedValue(undefined),
+    messageId: vi.fn().mockImplementation(() => activeMessageId),
+    previewMode: vi.fn().mockReturnValue("message"),
+    previewRevision: vi.fn().mockImplementation(() => previewRevision),
+    clear: vi.fn().mockResolvedValue(undefined),
+    stop: vi.fn().mockResolvedValue(undefined),
+    forceNewMessage: vi.fn().mockImplementation(() => {
+      activeMessageId = undefined;
+    }),
+    setMessageId: (value: number | undefined) => {
+      activeMessageId = value;
+    },
+  };
+}

src/telegram/lane-delivery.test.ts

@@ -1,42 +1,26 @@
 import { describe, expect, it, vi } from "vitest";
 import type { ReplyPayload } from "../auto-reply/types.js";
+import { createTestDraftStream } from "./draft-stream.test-helpers.js";
 import { createLaneTextDeliverer, type DraftLaneState, type LaneName } from "./lane-delivery.js";
 
-type MockStreamState = {
-  stream: NonNullable<DraftLaneState["stream"]>;
-  setMessageId: (value: number | undefined) => void;
-};
-
-function createMockStream(initialMessageId?: number): MockStreamState {
-  let messageId = initialMessageId;
-  const stream = {
-    update: vi.fn(),
-    flush: vi.fn().mockResolvedValue(undefined),
-    messageId: vi.fn().mockImplementation(() => messageId),
-    clear: vi.fn().mockResolvedValue(undefined),
-    stop: vi.fn().mockResolvedValue(undefined),
-    forceNewMessage: vi.fn(),
-    previewMode: vi.fn().mockReturnValue("message"),
-    previewRevision: vi.fn().mockReturnValue(0),
-  } as unknown as NonNullable<DraftLaneState["stream"]>;
-  return {
-    stream,
-    setMessageId: (value) => {
-      messageId = value;
-    },
-  };
-}
-
 function createHarness(params?: {
   answerMessageId?: number;
   draftMaxChars?: number;
   answerMessageIdAfterStop?: number;
 }) {
-  const answer = createMockStream(params?.answerMessageId);
+  const answer = createTestDraftStream({ messageId: params?.answerMessageId });
-  const reasoning = createMockStream();
+  const reasoning = createTestDraftStream();
   const lanes: Record<LaneName, DraftLaneState> = {
-    answer: { stream: answer.stream, lastPartialText: "", hasStreamedMessage: false },
+    answer: {
-    reasoning: { stream: reasoning.stream, lastPartialText: "", hasStreamedMessage: false },
+      stream: answer as DraftLaneState["stream"],
+      lastPartialText: "",
+      hasStreamedMessage: false,
+    },
+    reasoning: {
+      stream: reasoning as DraftLaneState["stream"],
+      lastPartialText: "",
+      hasStreamedMessage: false,
+    },
   };
   const sendPayload = vi.fn().mockResolvedValue(true);
   const flushDraftLane = vi.fn().mockImplementation(async (lane: DraftLaneState) => {
@@ -73,7 +57,10 @@ function createHarness(params?: {
   return {
     deliverLaneText,
     lanes,
-    answer,
+    answer: {
+      stream: answer,
+      setMessageId: answer.setMessageId,
+    },
     sendPayload,
     flushDraftLane,
     stopDraftLane,