From e58884925aee4ea6dd002018d18c4775b1143f73 Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Sun, 15 Feb 2026 22:12:07 +0000 Subject: [PATCH] refactor(test): reuse pi embedded subscribe session harness --- .../pi-embedded-subscribe.e2e-harness.ts | 10 +++++ ...esses-output-without-start-tag.e2e.test.ts | 44 +++++++------------ ...ion.subscribeembeddedpisession.e2e.test.ts | 22 ++++------ 3 files changed, 34 insertions(+), 42 deletions(-) diff --git a/src/agents/pi-embedded-subscribe.e2e-harness.ts b/src/agents/pi-embedded-subscribe.e2e-harness.ts index 88370693841..64975e8c72c 100644 --- a/src/agents/pi-embedded-subscribe.e2e-harness.ts +++ b/src/agents/pi-embedded-subscribe.e2e-harness.ts @@ -16,3 +16,13 @@ export function createStubSessionHarness(): { return { session, emit: (evt: unknown) => handler?.(evt) }; } + +export function extractAgentEventPayloads(calls: Array): Array> { + return calls + .map((call) => { + const first = call?.[0] as { data?: unknown } | undefined; + const data = first?.data; + return data && typeof data === "object" ? (data as Record) : undefined; + }) + .filter((value): value is Record => Boolean(value)); +} diff --git a/src/agents/pi-embedded-subscribe.subscribe-embedded-pi-session.filters-final-suppresses-output-without-start-tag.e2e.test.ts b/src/agents/pi-embedded-subscribe.subscribe-embedded-pi-session.filters-final-suppresses-output-without-start-tag.e2e.test.ts index 1dad92b6ce8..76a51a89197 100644 --- a/src/agents/pi-embedded-subscribe.subscribe-embedded-pi-session.filters-final-suppresses-output-without-start-tag.e2e.test.ts +++ b/src/agents/pi-embedded-subscribe.subscribe-embedded-pi-session.filters-final-suppresses-output-without-start-tag.e2e.test.ts @@ -1,34 +1,28 @@ import type { AssistantMessage } from "@mariozechner/pi-ai"; import { describe, expect, it, vi } from "vitest"; +import { + createStubSessionHarness, + extractAgentEventPayloads, +} from "./pi-embedded-subscribe.e2e-harness.js"; import { subscribeEmbeddedPiSession } from "./pi-embedded-subscribe.js"; -type StubSession = { - subscribe: (fn: (evt: unknown) => void) => () => void; -}; - describe("subscribeEmbeddedPiSession", () => { it("filters to and suppresses output without a start tag", () => { - let handler: ((evt: unknown) => void) | undefined; - const session: StubSession = { - subscribe: (fn) => { - handler = fn; - return () => {}; - }, - }; + const { session, emit } = createStubSessionHarness(); const onPartialReply = vi.fn(); const onAgentEvent = vi.fn(); subscribeEmbeddedPiSession({ - session: session as unknown as Parameters[0]["session"], + session, runId: "run", enforceFinalTag: true, onPartialReply, onAgentEvent, }); - handler?.({ type: "message_start", message: { role: "assistant" } }); - handler?.({ + emit({ type: "message_start", message: { role: "assistant" } }); + emit({ type: "message_update", message: { role: "assistant" }, assistantMessageEvent: { @@ -43,8 +37,8 @@ describe("subscribeEmbeddedPiSession", () => { onPartialReply.mockReset(); - handler?.({ type: "message_start", message: { role: "assistant" } }); - handler?.({ + emit({ type: "message_start", message: { role: "assistant" } }); + emit({ type: "message_update", message: { role: "assistant" }, assistantMessageEvent: { @@ -56,18 +50,12 @@ describe("subscribeEmbeddedPiSession", () => { expect(onPartialReply).not.toHaveBeenCalled(); }); it("emits agent events on message_end even without tags", () => { - let handler: ((evt: unknown) => void) | undefined; - const session: StubSession = { - subscribe: (fn) => { - handler = fn; - return () => {}; - }, - }; + const { session, emit } = createStubSessionHarness(); const onAgentEvent = vi.fn(); subscribeEmbeddedPiSession({ - session: session as unknown as Parameters[0]["session"], + session, runId: "run", enforceFinalTag: true, onAgentEvent, @@ -78,12 +66,10 @@ describe("subscribeEmbeddedPiSession", () => { content: [{ type: "text", text: "Hello world" }], } as AssistantMessage; - handler?.({ type: "message_start", message: assistantMessage }); - handler?.({ type: "message_end", message: assistantMessage }); + emit({ type: "message_start", message: assistantMessage }); + emit({ type: "message_end", message: assistantMessage }); - const payloads = onAgentEvent.mock.calls - .map((call) => call[0]?.data as Record | undefined) - .filter((value): value is Record => Boolean(value)); + const payloads = extractAgentEventPayloads(onAgentEvent.mock.calls); expect(payloads).toHaveLength(1); expect(payloads[0]?.text).toBe("Hello world"); expect(payloads[0]?.delta).toBe("Hello world"); diff --git a/src/agents/pi-embedded-subscribe.subscribe-embedded-pi-session.subscribeembeddedpisession.e2e.test.ts b/src/agents/pi-embedded-subscribe.subscribe-embedded-pi-session.subscribeembeddedpisession.e2e.test.ts index b53ffa62e53..1371a697d75 100644 --- a/src/agents/pi-embedded-subscribe.subscribe-embedded-pi-session.subscribeembeddedpisession.e2e.test.ts +++ b/src/agents/pi-embedded-subscribe.subscribe-embedded-pi-session.subscribeembeddedpisession.e2e.test.ts @@ -1,5 +1,9 @@ import type { AssistantMessage } from "@mariozechner/pi-ai"; import { describe, expect, it, vi } from "vitest"; +import { + createStubSessionHarness, + extractAgentEventPayloads, +} from "./pi-embedded-subscribe.e2e-harness.js"; import { subscribeEmbeddedPiSession } from "./pi-embedded-subscribe.js"; type StubSession = { @@ -186,18 +190,12 @@ describe("subscribeEmbeddedPiSession", () => { }); it("emits agent events on message_end for non-streaming assistant text", () => { - let handler: ((evt: unknown) => void) | undefined; - const session: StubSession = { - subscribe: (fn) => { - handler = fn; - return () => {}; - }, - }; + const { session, emit } = createStubSessionHarness(); const onAgentEvent = vi.fn(); subscribeEmbeddedPiSession({ - session: session as unknown as Parameters[0]["session"], + session, runId: "run", onAgentEvent, }); @@ -207,12 +205,10 @@ describe("subscribeEmbeddedPiSession", () => { content: [{ type: "text", text: "Hello world" }], } as AssistantMessage; - handler?.({ type: "message_start", message: assistantMessage }); - handler?.({ type: "message_end", message: assistantMessage }); + emit({ type: "message_start", message: assistantMessage }); + emit({ type: "message_end", message: assistantMessage }); - const payloads = onAgentEvent.mock.calls - .map((call) => call[0]?.data as Record | undefined) - .filter((value): value is Record => Boolean(value)); + const payloads = extractAgentEventPayloads(onAgentEvent.mock.calls); expect(payloads).toHaveLength(1); expect(payloads[0]?.text).toBe("Hello world"); expect(payloads[0]?.delta).toBe("Hello world");