From a52725e81a2718a46859e9cca19ae5d289915233 Mon Sep 17 00:00:00 2001
From: Vincent Koc <vincentkoc@ieee.org>
Date: Sun, 12 Apr 2026 10:58:35 +0100
Subject: [PATCH] test(agents): share sanitize history mock factories

---
 ...runner.openai-tool-id-preservation.test.ts | 22 +++++-----------
 ...er.sanitize-session-history.policy.test.ts | 26 +++++++------------
 ...r.sanitize-session-history.test-harness.ts | 23 ++++++++++++++++
 3 files changed, 40 insertions(+), 31 deletions(-)

diff --git a/src/agents/pi-embedded-runner.openai-tool-id-preservation.test.ts b/src/agents/pi-embedded-runner.openai-tool-id-preservation.test.ts
index a15d87f8c52..d42c9fd6944 100644
--- a/src/agents/pi-embedded-runner.openai-tool-id-preservation.test.ts
+++ b/src/agents/pi-embedded-runner.openai-tool-id-preservation.test.ts
@@ -1,6 +1,8 @@
 import type { AgentMessage } from "@mariozechner/pi-agent-core";
 import { beforeEach, describe, expect, it, vi } from "vitest";
 import {
+  createSanitizeSessionHistoryHelpersMock,
+  createSanitizeSessionHistoryProviderRuntimeMock,
   loadSanitizeSessionHistoryWithCleanMocks,
   makeInMemorySessionManager,
   makeModelSnapshotEntry,
@@ -8,22 +10,12 @@ import {
 } from "./pi-embedded-runner.sanitize-session-history.test-harness.js";
 import { castAgentMessage } from "./test-helpers/agent-message-fixtures.js";
 
-vi.mock("./pi-embedded-helpers.js", async () => ({
-  ...(await vi.importActual("./pi-embedded-helpers.js")),
-  sanitizeSessionMessagesImages: vi.fn(async (msgs) => msgs),
-}));
+vi.mock("./pi-embedded-helpers.js", async () => await createSanitizeSessionHistoryHelpersMock());
 
-vi.mock("../plugins/provider-runtime.js", async () => {
-  const actual = await vi.importActual<typeof import("../plugins/provider-runtime.js")>(
-    "../plugins/provider-runtime.js",
-  );
-  return {
-    ...actual,
-    resolveProviderRuntimePlugin: vi.fn(() => undefined),
-    sanitizeProviderReplayHistoryWithPlugin: vi.fn(() => undefined),
-    validateProviderReplayTurnsWithPlugin: vi.fn(() => undefined),
-  };
-});
+vi.mock(
+  "../plugins/provider-runtime.js",
+  async () => await createSanitizeSessionHistoryProviderRuntimeMock(),
+);
 
 describe("sanitizeSessionHistory openai tool id preservation", () => {
   let sanitizeSessionHistory: SanitizeSessionHistoryHarness["sanitizeSessionHistory"];
diff --git a/src/agents/pi-embedded-runner.sanitize-session-history.policy.test.ts b/src/agents/pi-embedded-runner.sanitize-session-history.policy.test.ts
index e691664d22c..3ff4d59801f 100644
--- a/src/agents/pi-embedded-runner.sanitize-session-history.policy.test.ts
+++ b/src/agents/pi-embedded-runner.sanitize-session-history.policy.test.ts
@@ -1,5 +1,7 @@
 import { beforeAll, beforeEach, describe, expect, it, vi } from "vitest";
 import {
+  createSanitizeSessionHistoryHelpersMock,
+  createSanitizeSessionHistoryProviderRuntimeMock,
   loadSanitizeSessionHistoryWithCleanMocks,
   makeMockSessionManager,
   makeSimpleUserMessages,
@@ -8,23 +10,15 @@ import {
   sanitizeWithOpenAIResponses,
 } from "./pi-embedded-runner.sanitize-session-history.test-harness.js";
 
-vi.mock("./pi-embedded-helpers.js", async () => ({
-  ...(await vi.importActual("./pi-embedded-helpers.js")),
-  isGoogleModelApi: vi.fn(),
-  sanitizeSessionMessagesImages: vi.fn(async (msgs) => msgs),
-}));
+vi.mock(
+  "./pi-embedded-helpers.js",
+  async () => await createSanitizeSessionHistoryHelpersMock({ isGoogleModelApi: vi.fn() }),
+);
 
-vi.mock("../plugins/provider-runtime.js", async () => {
-  const actual = await vi.importActual<typeof import("../plugins/provider-runtime.js")>(
-    "../plugins/provider-runtime.js",
-  );
-  return {
-    ...actual,
-    resolveProviderRuntimePlugin: vi.fn(() => undefined),
-    sanitizeProviderReplayHistoryWithPlugin: vi.fn(() => undefined),
-    validateProviderReplayTurnsWithPlugin: vi.fn(() => undefined),
-  };
-});
+vi.mock(
+  "../plugins/provider-runtime.js",
+  async () => await createSanitizeSessionHistoryProviderRuntimeMock(),
+);
 
 let sanitizeSessionHistory: SanitizeSessionHistoryHarness["sanitizeSessionHistory"];
 let mockedHelpers: SanitizeSessionHistoryHarness["mockedHelpers"];
diff --git a/src/agents/pi-embedded-runner.sanitize-session-history.test-harness.ts b/src/agents/pi-embedded-runner.sanitize-session-history.test-harness.ts
index 42b008279f5..faed4127eac 100644
--- a/src/agents/pi-embedded-runner.sanitize-session-history.test-harness.ts
+++ b/src/agents/pi-embedded-runner.sanitize-session-history.test-harness.ts
@@ -60,6 +60,29 @@ export function makeSimpleUserMessages(): AgentMessage[] {
   return messages as unknown as AgentMessage[];
 }
 
+export async function createSanitizeSessionHistoryHelpersMock(extra: Record<string, unknown> = {}) {
+  return {
+    ...(await vi.importActual("./pi-embedded-helpers.js")),
+    sanitizeSessionMessagesImages: vi.fn(async (msgs) => msgs),
+    ...extra,
+  };
+}
+
+export async function createSanitizeSessionHistoryProviderRuntimeMock(
+  extra: Record<string, unknown> = {},
+) {
+  const actual = await vi.importActual<typeof import("../plugins/provider-runtime.js")>(
+    "../plugins/provider-runtime.js",
+  );
+  return {
+    ...actual,
+    resolveProviderRuntimePlugin: vi.fn(() => undefined),
+    sanitizeProviderReplayHistoryWithPlugin: vi.fn(() => undefined),
+    validateProviderReplayTurnsWithPlugin: vi.fn(() => undefined),
+    ...extra,
+  };
+}
+
 export async function loadSanitizeSessionHistoryWithCleanMocks(): Promise<SanitizeSessionHistoryHarness> {
   vi.resetModules();
   vi.resetAllMocks();