From fdb0bf804f00951ef5cc155e3d4cf575e38e3469 Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Tue, 3 Mar 2026 02:42:14 +0000 Subject: [PATCH] refactor(test): dedupe telegram draft-stream fixtures --- src/telegram/bot-message-dispatch.test.ts | 45 +++----------- src/telegram/draft-stream.test-helpers.ts | 74 +++++++++++++++++++++++ src/telegram/lane-delivery.test.ts | 47 ++++++-------- 3 files changed, 98 insertions(+), 68 deletions(-) create mode 100644 src/telegram/draft-stream.test-helpers.ts diff --git a/src/telegram/bot-message-dispatch.test.ts b/src/telegram/bot-message-dispatch.test.ts index 39c04892216..ac79d0dc3c4 100644 --- a/src/telegram/bot-message-dispatch.test.ts +++ b/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) { - let previewRevision = 0; - return { - 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; - }), - }; - } + const createDraftStream = (messageId?: number) => createTestDraftStream({ messageId }); + const createSequencedDraftStream = (startMessageId = 1001) => + createSequencedTestDraftStream(startMessageId); function setupDraftStreams(params?: { answerMessageId?: number; reasoningMessageId?: number }) { const answerDraftStream = createDraftStream(params?.answerMessageId); diff --git a/src/telegram/draft-stream.test-helpers.ts b/src/telegram/draft-stream.test-helpers.ts new file mode 100644 index 00000000000..abb958e36f7 --- /dev/null +++ b/src/telegram/draft-stream.test-helpers.ts @@ -0,0 +1,74 @@ +import { vi } from "vitest"; + +type DraftPreviewMode = "message" | "draft"; + +export type TestDraftStream = { + update: ReturnType void>>; + flush: ReturnType Promise>>; + messageId: ReturnType number | undefined>>; + previewMode: ReturnType DraftPreviewMode>>; + previewRevision: ReturnType number>>; + clear: ReturnType Promise>>; + stop: ReturnType Promise>>; + forceNewMessage: ReturnType void>>; + setMessageId: (value: number | undefined) => void; +}; + +export function createTestDraftStream(params?: { + messageId?: number; + previewMode?: DraftPreviewMode; + onUpdate?: (text: string) => void; + onStop?: () => void | Promise; + 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; + }, + }; +} diff --git a/src/telegram/lane-delivery.test.ts b/src/telegram/lane-delivery.test.ts index f3599f0fde6..155fa7b63eb 100644 --- a/src/telegram/lane-delivery.test.ts +++ b/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; - 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; - return { - stream, - setMessageId: (value) => { - messageId = value; - }, - }; -} - function createHarness(params?: { answerMessageId?: number; draftMaxChars?: number; answerMessageIdAfterStop?: number; }) { - const answer = createMockStream(params?.answerMessageId); - const reasoning = createMockStream(); + const answer = createTestDraftStream({ messageId: params?.answerMessageId }); + const reasoning = createTestDraftStream(); const lanes: Record = { - answer: { stream: answer.stream, lastPartialText: "", hasStreamedMessage: false }, - reasoning: { stream: reasoning.stream, lastPartialText: "", hasStreamedMessage: false }, + answer: { + 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,