From e4287e0938e80a0ca410380edffaaba373d0f2b1 Mon Sep 17 00:00:00 2001
From: Peter Steinberger <steipete@gmail.com>
Date: Tue, 17 Mar 2026 07:27:52 +0000
Subject: [PATCH] refactor(compaction-tests): share snapshot assertions

---
 .../run/compaction-timeout.test.ts            | 40 +++++++++++++++----
 1 file changed, 32 insertions(+), 8 deletions(-)

diff --git a/src/agents/pi-embedded-runner/run/compaction-timeout.test.ts b/src/agents/pi-embedded-runner/run/compaction-timeout.test.ts
index 3853e0ebd25..54d6320297c 100644
--- a/src/agents/pi-embedded-runner/run/compaction-timeout.test.ts
+++ b/src/agents/pi-embedded-runner/run/compaction-timeout.test.ts
@@ -7,6 +7,30 @@ import {
   shouldFlagCompactionTimeout,
 } from "./compaction-timeout.js";
 
+function expectSelectedSnapshot(params: {
+  currentSessionId: string;
+  currentSnapshot: Parameters<typeof selectCompactionTimeoutSnapshot>[0]["currentSnapshot"];
+  expectedSessionIdUsed: string;
+  expectedSnapshot: ReadonlyArray<ReturnType<typeof castAgentMessage>>;
+  expectedSource: "current" | "pre-compaction";
+  preCompactionSessionId: string;
+  preCompactionSnapshot: Parameters<
+    typeof selectCompactionTimeoutSnapshot
+  >[0]["preCompactionSnapshot"];
+  timedOutDuringCompaction: boolean;
+}) {
+  const selected = selectCompactionTimeoutSnapshot({
+    timedOutDuringCompaction: params.timedOutDuringCompaction,
+    preCompactionSnapshot: params.preCompactionSnapshot,
+    preCompactionSessionId: params.preCompactionSessionId,
+    currentSnapshot: params.currentSnapshot,
+    currentSessionId: params.currentSessionId,
+  });
+  expect(selected.source).toBe(params.expectedSource);
+  expect(selected.sessionIdUsed).toBe(params.expectedSessionIdUsed);
+  expect(selected.messagesSnapshot).toEqual(params.expectedSnapshot);
+}
+
 describe("compaction-timeout helpers", () => {
   it("flags compaction timeout consistently for internal and external timeout sources", () => {
     const internalTimer = shouldFlagCompactionTimeout({
@@ -75,29 +99,29 @@ describe("compaction-timeout helpers", () => {
   it("uses pre-compaction snapshot when compaction timeout occurs", () => {
     const pre = [castAgentMessage({ role: "assistant", content: "pre" })] as const;
     const current = [castAgentMessage({ role: "assistant", content: "current" })] as const;
-    const selected = selectCompactionTimeoutSnapshot({
+    expectSelectedSnapshot({
       timedOutDuringCompaction: true,
       preCompactionSnapshot: [...pre],
       preCompactionSessionId: "session-pre",
       currentSnapshot: [...current],
       currentSessionId: "session-current",
+      expectedSource: "pre-compaction",
+      expectedSessionIdUsed: "session-pre",
+      expectedSnapshot: pre,
     });
-    expect(selected.source).toBe("pre-compaction");
-    expect(selected.sessionIdUsed).toBe("session-pre");
-    expect(selected.messagesSnapshot).toEqual(pre);
   });
 
   it("falls back to current snapshot when pre-compaction snapshot is unavailable", () => {
     const current = [castAgentMessage({ role: "assistant", content: "current" })] as const;
-    const selected = selectCompactionTimeoutSnapshot({
+    expectSelectedSnapshot({
       timedOutDuringCompaction: true,
       preCompactionSnapshot: null,
       preCompactionSessionId: "session-pre",
       currentSnapshot: [...current],
       currentSessionId: "session-current",
+      expectedSource: "current",
+      expectedSessionIdUsed: "session-current",
+      expectedSnapshot: current,
     });
-    expect(selected.source).toBe("current");
-    expect(selected.sessionIdUsed).toBe("session-current");
-    expect(selected.messagesSnapshot).toEqual(current);
   });
 });