From 03755f8463a5d59f1141c769ca7983dd36a1bb09 Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Tue, 3 Mar 2026 02:13:55 +0000 Subject: [PATCH] test(telegram): dedupe streaming cases and tighten sequential key checks --- src/telegram/bot-message-dispatch.test.ts | 104 +++-------- src/telegram/bot.create-telegram-bot.test.ts | 175 +++++++++---------- 2 files changed, 109 insertions(+), 170 deletions(-) diff --git a/src/telegram/bot-message-dispatch.test.ts b/src/telegram/bot-message-dispatch.test.ts index 5104c7c053d..66f1fb33e19 100644 --- a/src/telegram/bot-message-dispatch.test.ts +++ b/src/telegram/bot-message-dispatch.test.ts @@ -588,7 +588,10 @@ describe("dispatchTelegramMessage draft streaming", () => { expect(draftStream.stop).toHaveBeenCalled(); }); - it("disables block streaming when streamMode is off", async () => { + it.each([ + { label: "default account config", telegramCfg: {} }, + { label: "account blockStreaming override", telegramCfg: { blockStreaming: true } }, + ])("disables block streaming when streamMode is off ($label)", async ({ telegramCfg }) => { dispatchReplyWithBufferedBlockDispatcher.mockImplementation(async ({ dispatcherOptions }) => { await dispatcherOptions.deliver({ text: "Hello" }, { kind: "final" }); return { queuedFinal: true }; @@ -598,6 +601,7 @@ describe("dispatchTelegramMessage draft streaming", () => { await dispatchWithContext({ context: createContext(), streamMode: "off", + telegramCfg, }); expect(createTelegramDraftStream).not.toHaveBeenCalled(); @@ -610,69 +614,27 @@ describe("dispatchTelegramMessage draft streaming", () => { ); }); - it("disables block streaming when streamMode is off even if blockStreaming config is true", async () => { - dispatchReplyWithBufferedBlockDispatcher.mockImplementation(async ({ dispatcherOptions }) => { - await dispatcherOptions.deliver({ text: "Hello" }, { kind: "final" }); - return { queuedFinal: true }; - }); - deliverReplies.mockResolvedValue({ delivered: true }); + it.each(["block", "partial"] as const)( + "forces new message when assistant message restarts (%s mode)", + async (streamMode) => { + const draftStream = createDraftStream(999); + createTelegramDraftStream.mockReturnValue(draftStream); + dispatchReplyWithBufferedBlockDispatcher.mockImplementation( + async ({ dispatcherOptions, replyOptions }) => { + await replyOptions?.onPartialReply?.({ text: "First response" }); + await replyOptions?.onAssistantMessageStart?.(); + await replyOptions?.onPartialReply?.({ text: "After tool call" }); + await dispatcherOptions.deliver({ text: "After tool call" }, { kind: "final" }); + return { queuedFinal: true }; + }, + ); + deliverReplies.mockResolvedValue({ delivered: true }); - await dispatchWithContext({ - context: createContext(), - streamMode: "off", - telegramCfg: { blockStreaming: true }, - }); + await dispatchWithContext({ context: createContext(), streamMode }); - expect(createTelegramDraftStream).not.toHaveBeenCalled(); - expect(dispatchReplyWithBufferedBlockDispatcher).toHaveBeenCalledWith( - expect.objectContaining({ - replyOptions: expect.objectContaining({ - disableBlockStreaming: true, - }), - }), - ); - }); - - it("forces new message for next assistant block in legacy block stream mode", async () => { - const draftStream = createDraftStream(999); - createTelegramDraftStream.mockReturnValue(draftStream); - dispatchReplyWithBufferedBlockDispatcher.mockImplementation( - async ({ dispatcherOptions, replyOptions }) => { - // First assistant message: partial text - await replyOptions?.onPartialReply?.({ text: "First response" }); - // New assistant message starts (e.g., after tool call) - await replyOptions?.onAssistantMessageStart?.(); - // Second assistant message: new text - await replyOptions?.onPartialReply?.({ text: "After tool call" }); - await dispatcherOptions.deliver({ text: "After tool call" }, { kind: "final" }); - return { queuedFinal: true }; - }, - ); - deliverReplies.mockResolvedValue({ delivered: true }); - - await dispatchWithContext({ context: createContext(), streamMode: "block" }); - - expect(draftStream.forceNewMessage).toHaveBeenCalledTimes(1); - }); - - it("forces new message in partial mode when assistant message restarts", async () => { - const draftStream = createDraftStream(999); - createTelegramDraftStream.mockReturnValue(draftStream); - dispatchReplyWithBufferedBlockDispatcher.mockImplementation( - async ({ dispatcherOptions, replyOptions }) => { - await replyOptions?.onPartialReply?.({ text: "First response" }); - await replyOptions?.onAssistantMessageStart?.(); - await replyOptions?.onPartialReply?.({ text: "After tool call" }); - await dispatcherOptions.deliver({ text: "After tool call" }, { kind: "final" }); - return { queuedFinal: true }; - }, - ); - deliverReplies.mockResolvedValue({ delivered: true }); - - await dispatchWithContext({ context: createContext(), streamMode: "partial" }); - - expect(draftStream.forceNewMessage).toHaveBeenCalledTimes(1); - }); + expect(draftStream.forceNewMessage).toHaveBeenCalledTimes(1); + }, + ); it("does not force new message on first assistant message start", async () => { const draftStream = createDraftStream(999); @@ -1076,7 +1038,7 @@ describe("dispatchTelegramMessage draft streaming", () => { it.each([undefined, null] as const)( "skips outbound send when final payload text is %s and has no media", async (emptyText) => { - setupDraftStreams({ answerMessageId: 999 }); + const { answerDraftStream } = setupDraftStreams({ answerMessageId: 999 }); dispatchReplyWithBufferedBlockDispatcher.mockImplementation(async ({ dispatcherOptions }) => { await dispatcherOptions.deliver( { text: emptyText as unknown as string }, @@ -1090,6 +1052,7 @@ describe("dispatchTelegramMessage draft streaming", () => { expect(deliverReplies).not.toHaveBeenCalled(); expect(editMessageTelegram).not.toHaveBeenCalled(); + expect(answerDraftStream.clear).toHaveBeenCalledTimes(1); }, ); @@ -1595,21 +1558,6 @@ describe("dispatchTelegramMessage draft streaming", () => { expect(draftStream.clear).toHaveBeenCalledTimes(1); }); - it("skips final payload when text is undefined", async () => { - const draftStream = createDraftStream(999); - createTelegramDraftStream.mockReturnValue(draftStream); - dispatchReplyWithBufferedBlockDispatcher.mockImplementation(async ({ dispatcherOptions }) => { - await dispatcherOptions.deliver({ text: undefined as unknown as string }, { kind: "final" }); - return { queuedFinal: true }; - }); - deliverReplies.mockResolvedValue({ delivered: true }); - - await dispatchWithContext({ context: createContext() }); - - expect(deliverReplies).not.toHaveBeenCalled(); - expect(draftStream.clear).toHaveBeenCalledTimes(1); - }); - it("falls back when all finals are skipped and clears preview", async () => { const draftStream = createDraftStream(999); createTelegramDraftStream.mockReturnValue(draftStream); diff --git a/src/telegram/bot.create-telegram-bot.test.ts b/src/telegram/bot.create-telegram-bot.test.ts index 50e104a401d..2b29bbe746a 100644 --- a/src/telegram/bot.create-telegram-bot.test.ts +++ b/src/telegram/bot.create-telegram-bot.test.ts @@ -5,6 +5,7 @@ import type { Chat, Message } from "@grammyjs/types"; import { afterAll, beforeAll, describe, expect, it, vi } from "vitest"; import { escapeRegExp, formatEnvelopeTimestamp } from "../../test/helpers/envelope-timestamp.js"; import { withEnvAsync } from "../test-utils/env.js"; +import { useFrozenTime, useRealTime } from "../test-utils/frozen-time.js"; import { answerCallbackQuerySpy, botCtorSpy, @@ -123,97 +124,87 @@ describe("createTelegramBot", () => { expect(sequentializeSpy).toHaveBeenCalledTimes(1); expect(middlewareUseSpy).toHaveBeenCalledWith(sequentializeSpy.mock.results[0]?.value); expect(sequentializeKey).toBe(getTelegramSequentialKey); - expect( - getTelegramSequentialKey({ message: mockMessage({ chat: mockChat({ id: 123 }) }) }), - ).toBe("telegram:123"); - expect( - getTelegramSequentialKey({ - message: mockMessage({ - chat: mockChat({ id: 123, type: "private" }), - message_thread_id: 9, - }), - }), - ).toBe("telegram:123:topic:9"); - expect( - getTelegramSequentialKey({ - message: mockMessage({ - chat: mockChat({ id: 123, type: "supergroup" }), - message_thread_id: 9, - }), - }), - ).toBe("telegram:123"); - expect( - getTelegramSequentialKey({ - message: mockMessage({ chat: mockChat({ id: 123, type: "supergroup", is_forum: true }) }), - }), - ).toBe("telegram:123:topic:1"); - expect( - getTelegramSequentialKey({ - update: { message: mockMessage({ chat: mockChat({ id: 555 }) }) }, - }), - ).toBe("telegram:555"); - expect( - getTelegramSequentialKey({ - channelPost: mockMessage({ chat: mockChat({ id: -100777111222, type: "channel" }) }), - }), - ).toBe("telegram:-100777111222"); - expect( - getTelegramSequentialKey({ - update: { - channel_post: mockMessage({ chat: mockChat({ id: -100777111223, type: "channel" }) }), + const cases = [ + [{ message: mockMessage({ chat: mockChat({ id: 123 }) }) }, "telegram:123"], + [ + { + message: mockMessage({ + chat: mockChat({ id: 123, type: "private" }), + message_thread_id: 9, + }), }, - }), - ).toBe("telegram:-100777111223"); - expect( - getTelegramSequentialKey({ - message: mockMessage({ chat: mockChat({ id: 123 }), text: "/stop" }), - }), - ).toBe("telegram:123:control"); - expect( - getTelegramSequentialKey({ - message: mockMessage({ chat: mockChat({ id: 123 }), text: "/status" }), - }), - ).toBe("telegram:123"); - expect( - getTelegramSequentialKey({ - message: mockMessage({ chat: mockChat({ id: 123 }), text: "stop" }), - }), - ).toBe("telegram:123:control"); - expect( - getTelegramSequentialKey({ - message: mockMessage({ chat: mockChat({ id: 123 }), text: "stop please" }), - }), - ).toBe("telegram:123:control"); - expect( - getTelegramSequentialKey({ - message: mockMessage({ chat: mockChat({ id: 123 }), text: "do not do that" }), - }), - ).toBe("telegram:123:control"); - expect( - getTelegramSequentialKey({ - message: mockMessage({ chat: mockChat({ id: 123 }), text: "остановись" }), - }), - ).toBe("telegram:123:control"); - expect( - getTelegramSequentialKey({ - message: mockMessage({ chat: mockChat({ id: 123 }), text: "halt" }), - }), - ).toBe("telegram:123:control"); - expect( - getTelegramSequentialKey({ - message: mockMessage({ chat: mockChat({ id: 123 }), text: "/abort" }), - }), - ).toBe("telegram:123"); - expect( - getTelegramSequentialKey({ - message: mockMessage({ chat: mockChat({ id: 123 }), text: "/abort now" }), - }), - ).toBe("telegram:123"); - expect( - getTelegramSequentialKey({ - message: mockMessage({ chat: mockChat({ id: 123 }), text: "please do not do that" }), - }), - ).toBe("telegram:123"); + "telegram:123:topic:9", + ], + [ + { + message: mockMessage({ + chat: mockChat({ id: 123, type: "supergroup" }), + message_thread_id: 9, + }), + }, + "telegram:123", + ], + [ + { + message: mockMessage({ + chat: mockChat({ id: 123, type: "supergroup", is_forum: true }), + }), + }, + "telegram:123:topic:1", + ], + [{ update: { message: mockMessage({ chat: mockChat({ id: 555 }) }) } }, "telegram:555"], + [ + { + channelPost: mockMessage({ chat: mockChat({ id: -100777111222, type: "channel" }) }), + }, + "telegram:-100777111222", + ], + [ + { + update: { + channel_post: mockMessage({ chat: mockChat({ id: -100777111223, type: "channel" }) }), + }, + }, + "telegram:-100777111223", + ], + [ + { message: mockMessage({ chat: mockChat({ id: 123 }), text: "/stop" }) }, + "telegram:123:control", + ], + [{ message: mockMessage({ chat: mockChat({ id: 123 }), text: "/status" }) }, "telegram:123"], + [ + { message: mockMessage({ chat: mockChat({ id: 123 }), text: "stop" }) }, + "telegram:123:control", + ], + [ + { message: mockMessage({ chat: mockChat({ id: 123 }), text: "stop please" }) }, + "telegram:123:control", + ], + [ + { message: mockMessage({ chat: mockChat({ id: 123 }), text: "do not do that" }) }, + "telegram:123:control", + ], + [ + { message: mockMessage({ chat: mockChat({ id: 123 }), text: "остановись" }) }, + "telegram:123:control", + ], + [ + { message: mockMessage({ chat: mockChat({ id: 123 }), text: "halt" }) }, + "telegram:123:control", + ], + [{ message: mockMessage({ chat: mockChat({ id: 123 }), text: "/abort" }) }, "telegram:123"], + [ + { message: mockMessage({ chat: mockChat({ id: 123 }), text: "/abort now" }) }, + "telegram:123", + ], + [ + { message: mockMessage({ chat: mockChat({ id: 123 }), text: "please do not do that" }) }, + "telegram:123", + ], + ] as const; + for (const [input, expected] of cases) { + expect(getTelegramSequentialKey(input)).toBe(expected); + } }); it("routes callback_query payloads as messages and answers callbacks", async () => { createTelegramBot({ token: "tok" }); @@ -2031,7 +2022,7 @@ describe("createTelegramBot", () => { }, }); - vi.useFakeTimers(); + useFrozenTime("2026-02-20T00:00:00.000Z"); try { createTelegramBot({ token: "tok", testTimings: TELEGRAM_TEST_TIMINGS }); const handler = getOnHandler("channel_post") as ( @@ -2071,7 +2062,7 @@ describe("createTelegramBot", () => { expect(payload.RawBody).toContain(part1.slice(0, 32)); expect(payload.RawBody).toContain(part2.slice(0, 32)); } finally { - vi.useRealTimers(); + useRealTime(); } }); it("drops oversized channel_post media instead of dispatching a placeholder message", async () => {