test: speed up embedded runner e2e mocks

2026-05-06 10:10:45 +00:00 · 2026-04-27 11:50:33 +01:00
parent 836d4b4105
commit c9b9887583
4 changed files with 174 additions and 168 deletions
--- a/src/agents/model-fallback.run-embedded.e2e.test.ts
+++ b/src/agents/model-fallback.run-embedded.e2e.test.ts
@@ -12,7 +12,11 @@ import {
  createResolvedEmbeddedRunnerModel,
  makeEmbeddedRunnerAttempt,
 } from "./test-helpers/pi-embedded-runner-e2e-fixtures.js";
-import { installEmbeddedRunnerBaseE2eMocks } from "./test-helpers/pi-embedded-runner-e2e-mocks.js";
+import {
+  installEmbeddedRunnerBackoffE2eMocks,
+  installEmbeddedRunnerBaseE2eMocks,
+  installEmbeddedRunnerFastRunE2eMocks,
+} from "./test-helpers/pi-embedded-runner-e2e-mocks.js";

 const runEmbeddedAttemptMock = vi.fn<(params: unknown) => Promise<EmbeddedRunAttemptResult>>();
 const { computeBackoffMock, sleepWithAbortMock } = vi.hoisted(() => ({
@@ -25,28 +29,6 @@ const { computeBackoffMock, sleepWithAbortMock } = vi.hoisted(() => ({
  sleepWithAbortMock: vi.fn(async (_ms: number, _abortSignal?: AbortSignal) => undefined),
 }));

-vi.mock("./pi-embedded-runner/run/attempt.js", async () => {
-  const actual = await vi.importActual<typeof import("./pi-embedded-runner/run/attempt.js")>(
-    "./pi-embedded-runner/run/attempt.js",
-  );
-  return {
-    ...actual,
-    runEmbeddedAttempt: (params: unknown) => runEmbeddedAttemptMock(params),
-  };
-});
-
-vi.mock("../infra/backoff.js", async () => {
-  const actual = await vi.importActual<typeof import("../infra/backoff.js")>("../infra/backoff.js");
-  return {
-    ...actual,
-    computeBackoff: (
-      policy: { initialMs: number; maxMs: number; factor: number; jitter: number },
-      attempt: number,
-    ) => computeBackoffMock(policy, attempt),
-    sleepWithAbort: (ms: number, abortSignal?: AbortSignal) => sleepWithAbortMock(ms, abortSignal),
-  };
-});
-
 vi.mock("./models-config.js", async () => {
  const mod = await vi.importActual<typeof import("./models-config.js")>("./models-config.js");
  return {
@@ -57,20 +39,17 @@ vi.mock("./models-config.js", async () => {

 const installRunEmbeddedMocks = () => {
  installEmbeddedRunnerBaseE2eMocks();
+  installEmbeddedRunnerFastRunE2eMocks({
+    runEmbeddedAttempt: (params) => runEmbeddedAttemptMock(params),
+  });
+  installEmbeddedRunnerBackoffE2eMocks({
+    computeBackoff: (policy, attempt) => computeBackoffMock(policy, attempt),
+    sleepWithAbort: (ms, abortSignal) => sleepWithAbortMock(ms, abortSignal),
+  });
  vi.doMock("./pi-embedded-runner/model.js", () => ({
    resolveModelAsync: async (provider: string, modelId: string) =>
      createResolvedEmbeddedRunnerModel(provider, modelId),
  }));
-  vi.doMock("../plugins/provider-runtime.js", async () => {
-    const actual = await vi.importActual<typeof import("../plugins/provider-runtime.js")>(
-      "../plugins/provider-runtime.js",
-    );
-    return {
-      ...actual,
-      prepareProviderRuntimeAuth: vi.fn(async () => undefined),
-      resolveProviderCapabilitiesWithPlugin: vi.fn(() => undefined),
-    };
-  });
 };

 let runEmbeddedPiAgent: typeof import("./pi-embedded-runner/run.js").runEmbeddedPiAgent;
--- a/src/agents/pi-embedded-runner.e2e.test.ts
+++ b/src/agents/pi-embedded-runner.e2e.test.ts
@@ -13,7 +13,10 @@ import {
  immediateEnqueue,
  makeEmbeddedRunnerAttempt,
 } from "./test-helpers/pi-embedded-runner-e2e-fixtures.js";
-import { installEmbeddedRunnerBaseE2eMocks } from "./test-helpers/pi-embedded-runner-e2e-mocks.js";
+import {
+  installEmbeddedRunnerBaseE2eMocks,
+  installEmbeddedRunnerFastRunE2eMocks,
+} from "./test-helpers/pi-embedded-runner-e2e-mocks.js";

 const runEmbeddedAttemptMock = vi.fn();
 const disposeSessionMcpRuntimeMock = vi.fn<(sessionId: string) => Promise<void>>(async () => {
@@ -88,6 +91,9 @@ vi.mock("@mariozechner/pi-ai", async () => {

 const installRunEmbeddedMocks = () => {
  installEmbeddedRunnerBaseE2eMocks({ hookRunner: "full" });
+  installEmbeddedRunnerFastRunE2eMocks({
+    runEmbeddedAttempt: (params) => runEmbeddedAttemptMock(params),
+  });
  vi.doMock("./command/session.js", async () => {
    const actual =
      await vi.importActual<typeof import("./command/session.js")>("./command/session.js");
@@ -110,9 +116,6 @@ const installRunEmbeddedMocks = () => {
      },
    };
  });
-  vi.doMock("./pi-embedded-runner/run/attempt.js", () => ({
-    runEmbeddedAttempt: (params: unknown) => runEmbeddedAttemptMock(params),
-  }));
  vi.doMock("./pi-bundle-mcp-tools.js", () => ({
    disposeSessionMcpRuntime: (sessionId: string) => disposeSessionMcpRuntimeMock(sessionId),
    retireSessionMcpRuntimeForSessionKey: () => Promise.resolve(false),
@@ -139,15 +142,6 @@ const installRunEmbeddedMocks = () => {
      stopRuntimeAuthRefreshTimer: vi.fn(),
    }),
  }));
-  vi.doMock("../plugins/provider-runtime.js", async () => {
-    const actual = await vi.importActual<typeof import("../plugins/provider-runtime.js")>(
-      "../plugins/provider-runtime.js",
-    );
-    return {
-      ...actual,
-      prepareProviderRuntimeAuth: vi.fn(async () => undefined),
-    };
-  });
  vi.doMock("./models-config.js", async () => {
    const mod = await vi.importActual<typeof import("./models-config.js")>("./models-config.js");
    return {
--- a/src/agents/pi-embedded-runner.run-embedded-pi-agent.auth-profile-rotation.e2e.test.ts
+++ b/src/agents/pi-embedded-runner.run-embedded-pi-agent.auth-profile-rotation.e2e.test.ts
@@ -8,6 +8,11 @@ import { redactIdentifier } from "../logging/redact-identifier.js";
 import type { AuthProfileFailureReason } from "./auth-profiles.js";
 import { buildAttemptReplayMetadata } from "./pi-embedded-runner/run/incomplete-turn.js";
 import type { EmbeddedRunAttemptResult } from "./pi-embedded-runner/run/types.js";
+import {
+  installEmbeddedRunnerBackoffE2eMocks,
+  installEmbeddedRunnerBaseE2eMocks,
+  installEmbeddedRunnerFastRunE2eMocks,
+} from "./test-helpers/pi-embedded-runner-e2e-mocks.js";

 const runEmbeddedAttemptMock = vi.fn<(params: unknown) => Promise<EmbeddedRunAttemptResult>>();
 const resolveCopilotApiTokenMock = vi.fn();
@@ -22,96 +27,21 @@ const { computeBackoffMock, sleepWithAbortMock } = vi.hoisted(() => ({
 }));

 const installRunEmbeddedMocks = () => {
-  vi.doMock("../plugins/hook-runner-global.js", () => ({
-    getGlobalHookRunner: vi.fn(() => undefined),
-  }));
-  vi.doMock("../context-engine/init.js", () => ({
-    ensureContextEnginesInitialized: vi.fn(),
-  }));
-  vi.doMock("../context-engine/registry.js", () => ({
-    resolveContextEngine: vi.fn(async () => ({
-      dispose: async () => undefined,
-    })),
-  }));
-  vi.doMock("./runtime-plugins.js", () => ({
-    ensureRuntimePluginsLoaded: vi.fn(),
-  }));
-  vi.doMock("./harness/selection.js", () => ({
-    selectAgentHarness: vi.fn((params: { provider?: string }) => ({
-      id: params.provider === "codex-cli" ? "codex" : "pi",
-      label: "Mock agent harness",
-      supports: vi.fn(() => ({ supported: false })),
-      runAttempt: vi.fn(),
-    })),
-    runAgentHarnessAttemptWithFallback: (params: unknown) => runEmbeddedAttemptMock(params),
-  }));
-  vi.doMock("./runtime-plan/build.js", () => ({
-    buildAgentRuntimePlan: vi.fn(
-      (params: {
-        provider: string;
-        modelId: string;
-        modelApi?: string | null;
-        harnessId?: string;
-        sessionAuthProfileId?: string;
-      }) => ({
-        resolvedRef: {
-          provider: params.provider,
-          modelId: params.modelId,
-          ...(params.modelApi ? { modelApi: params.modelApi } : {}),
-          ...(params.harnessId ? { harnessId: params.harnessId } : {}),
-        },
-        auth: {
-          providerForAuth: params.provider,
-          authProfileProviderForAuth: params.sessionAuthProfileId?.split(":", 1)[0] ?? "",
-          forwardedAuthProfileId: params.sessionAuthProfileId,
-        },
-        prompt: {
-          provider: params.provider,
-          modelId: params.modelId,
-          resolveSystemPromptContribution: vi.fn(() => undefined),
-        },
-        tools: {
-          normalize: vi.fn((tools: unknown[]) => tools),
-          logDiagnostics: vi.fn(),
-        },
-        transcript: {
-          policy: {
-            sanitizeMode: "full",
-            sanitizeToolCallIds: true,
-            preserveNativeAnthropicToolUseIds: false,
-            repairToolUseResultPairing: true,
-            preserveSignatures: false,
-            sanitizeThinkingSignatures: true,
-            dropThinkingBlocks: false,
-            applyGoogleTurnOrdering: false,
-            validateGeminiTurns: false,
-            validateAnthropicTurns: false,
-            allowSyntheticToolResults: true,
-          },
-          resolvePolicy: vi.fn(() => undefined),
-        },
-        delivery: {
-          isSilentPayload: vi.fn(() => false),
-          resolveFollowupRoute: vi.fn(() => undefined),
-        },
-        outcome: {
-          classifyRunResult: vi.fn(() => undefined),
-        },
-        transport: {
-          extraParams: {},
-          resolveExtraParams: vi.fn(() => ({})),
-        },
-        observability: {
-          resolvedRef: `${params.provider}/${params.modelId}`,
-          provider: params.provider,
-          modelId: params.modelId,
-          ...(params.modelApi ? { modelApi: params.modelApi } : {}),
-          ...(params.harnessId ? { harnessId: params.harnessId } : {}),
-          ...(params.sessionAuthProfileId ? { authProfileId: params.sessionAuthProfileId } : {}),
-        },
-      }),
-    ),
-  }));
+  installEmbeddedRunnerBaseE2eMocks();
+  installEmbeddedRunnerFastRunE2eMocks({
+    runEmbeddedAttempt: (params) => runEmbeddedAttemptMock(params),
+    prepareProviderRuntimeAuth: async (params) => {
+      if (params.provider !== "github-copilot") {
+        return undefined;
+      }
+      const token = await resolveCopilotApiTokenMock(params.context.apiKey);
+      return {
+        apiKey: token.token,
+        baseUrl: token.baseUrl,
+        expiresAt: token.expiresAt,
+      };
+    },
+  });
  vi.doMock("./pi-embedded-runner/model.js", () => ({
    resolveModelAsync: async (provider: string, modelId: string) => ({
      model: {
@@ -134,38 +64,10 @@ const installRunEmbeddedMocks = () => {
      modelRegistry: {},
    }),
  }));
-  vi.doMock("./pi-embedded-runner/run/attempt.js", () => ({
-    runEmbeddedAttempt: (params: unknown) => runEmbeddedAttemptMock(params),
-  }));
-  vi.doMock("../plugins/provider-runtime.js", () => ({
-    buildProviderMissingAuthMessageWithPlugin: vi.fn(() => undefined),
-    prepareProviderRuntimeAuth: async (params: {
-      provider: string;
-      context: { apiKey: string };
-    }) => {
-      if (params.provider !== "github-copilot") {
-        return undefined;
-      }
-      const token = await resolveCopilotApiTokenMock(params.context.apiKey);
-      return {
-        apiKey: token.token,
-        baseUrl: token.baseUrl,
-        expiresAt: token.expiresAt,
-      };
-    },
-    resolveProviderAuthProfileId: vi.fn(() => undefined),
-    resolveProviderCapabilitiesWithPlugin: vi.fn(() => undefined),
-    resolveExternalAuthProfilesWithPlugins: vi.fn(() => []),
-    resolveProviderSyntheticAuthWithPlugin: vi.fn(() => undefined),
-    shouldDeferProviderSyntheticProfileAuthWithPlugin: vi.fn(() => false),
-  }));
-  vi.doMock("../infra/backoff.js", () => ({
-    computeBackoff: (
-      policy: { initialMs: number; maxMs: number; factor: number; jitter: number },
-      attempt: number,
-    ) => computeBackoffMock(policy, attempt),
-    sleepWithAbort: (ms: number, abortSignal?: AbortSignal) => sleepWithAbortMock(ms, abortSignal),
-  }));
+  installEmbeddedRunnerBackoffE2eMocks({
+    computeBackoff: (policy, attempt) => computeBackoffMock(policy, attempt),
+    sleepWithAbort: (ms, abortSignal) => sleepWithAbortMock(ms, abortSignal),
+  });
  vi.doMock("./pi-embedded-runner/compact.js", () => ({
    compactEmbeddedPiSessionDirect: vi.fn(async () => {
      throw new Error("compact should not run in auth profile rotation tests");
--- a/src/agents/test-helpers/pi-embedded-runner-e2e-mocks.ts
+++ b/src/agents/test-helpers/pi-embedded-runner-e2e-mocks.ts
@@ -1,5 +1,21 @@
 import { vi } from "vitest";

+type EmbeddedRunnerFastRunMockOptions = {
+  runEmbeddedAttempt: (params: unknown) => unknown;
+  prepareProviderRuntimeAuth?: (params: {
+    provider: string;
+    context: { apiKey: string };
+  }) => unknown;
+};
+
+type EmbeddedRunnerBackoffMockOptions = {
+  computeBackoff: (
+    policy: { initialMs: number; maxMs: number; factor: number; jitter: number },
+    attempt: number,
+  ) => number;
+  sleepWithAbort: (ms: number, abortSignal?: AbortSignal) => unknown;
+};
+
 export function installEmbeddedRunnerBaseE2eMocks(options?: {
  hookRunner?: "minimal" | "full";
 }): void {
@@ -28,3 +44,118 @@ export function installEmbeddedRunnerBaseE2eMocks(options?: {
    ensureRuntimePluginsLoaded: vi.fn(),
  }));
 }
+
+export function installEmbeddedRunnerFastRunE2eMocks(
+  options: EmbeddedRunnerFastRunMockOptions,
+): void {
+  vi.doMock("../harness/selection.js", () => ({
+    selectAgentHarness: vi.fn((params: { provider?: string }) => ({
+      id: params.provider === "codex-cli" ? "codex" : "pi",
+      label: "Mock agent harness",
+      supports: vi.fn(() => ({ supported: false })),
+      runAttempt: vi.fn(),
+    })),
+    runAgentHarnessAttemptWithFallback: (params: unknown) => options.runEmbeddedAttempt(params),
+  }));
+  vi.doMock("../runtime-plan/build.js", () => ({
+    buildAgentRuntimePlan: vi.fn(
+      (params: {
+        provider: string;
+        modelId: string;
+        modelApi?: string | null;
+        harnessId?: string;
+        sessionAuthProfileId?: string;
+      }) => ({
+        resolvedRef: {
+          provider: params.provider,
+          modelId: params.modelId,
+          ...(params.modelApi ? { modelApi: params.modelApi } : {}),
+          ...(params.harnessId ? { harnessId: params.harnessId } : {}),
+        },
+        auth: {
+          providerForAuth: params.provider,
+          authProfileProviderForAuth: params.sessionAuthProfileId?.split(":", 1)[0] ?? "",
+          forwardedAuthProfileId: params.sessionAuthProfileId,
+        },
+        prompt: {
+          provider: params.provider,
+          modelId: params.modelId,
+          resolveSystemPromptContribution: vi.fn(() => undefined),
+        },
+        tools: {
+          normalize: vi.fn((tools: unknown[]) => tools),
+          logDiagnostics: vi.fn(),
+        },
+        transcript: {
+          policy: {
+            sanitizeMode: "full",
+            sanitizeToolCallIds: true,
+            preserveNativeAnthropicToolUseIds: false,
+            repairToolUseResultPairing: true,
+            preserveSignatures: false,
+            sanitizeThinkingSignatures: true,
+            dropThinkingBlocks: false,
+            applyGoogleTurnOrdering: false,
+            validateGeminiTurns: false,
+            validateAnthropicTurns: false,
+            allowSyntheticToolResults: true,
+          },
+          resolvePolicy: vi.fn(() => undefined),
+        },
+        delivery: {
+          isSilentPayload: vi.fn(() => false),
+          resolveFollowupRoute: vi.fn(() => undefined),
+        },
+        outcome: {
+          classifyRunResult: vi.fn(() => undefined),
+        },
+        transport: {
+          extraParams: {},
+          resolveExtraParams: vi.fn(() => ({})),
+        },
+        observability: {
+          resolvedRef: `${params.provider}/${params.modelId}`,
+          provider: params.provider,
+          modelId: params.modelId,
+          ...(params.modelApi ? { modelApi: params.modelApi } : {}),
+          ...(params.harnessId ? { harnessId: params.harnessId } : {}),
+          ...(params.sessionAuthProfileId ? { authProfileId: params.sessionAuthProfileId } : {}),
+        },
+      }),
+    ),
+  }));
+  vi.doMock("../pi-embedded-runner/run/attempt.js", () => ({
+    runEmbeddedAttempt: (params: unknown) => options.runEmbeddedAttempt(params),
+  }));
+  vi.doMock("../../plugins/provider-runtime.js", () => ({
+    applyProviderResolvedModelCompatWithPlugins: vi.fn(() => undefined),
+    applyProviderResolvedTransportWithPlugin: vi.fn(() => undefined),
+    buildProviderMissingAuthMessageWithPlugin: vi.fn(() => undefined),
+    buildProviderUnknownModelHintWithPlugin: vi.fn(() => undefined),
+    clearProviderRuntimeHookCache: vi.fn(),
+    normalizeProviderResolvedModelWithPlugin: vi.fn(() => undefined),
+    normalizeProviderTransportWithPlugin: vi.fn(() => undefined),
+    prepareProviderDynamicModel: vi.fn(async () => undefined),
+    prepareProviderRuntimeAuth: options.prepareProviderRuntimeAuth ?? vi.fn(async () => undefined),
+    resolveProviderAuthProfileId: vi.fn(() => undefined),
+    resolveProviderCapabilitiesWithPlugin: vi.fn(() => undefined),
+    resolveExternalAuthProfilesWithPlugins: vi.fn(() => []),
+    resolveProviderSyntheticAuthWithPlugin: vi.fn(() => undefined),
+    runProviderDynamicModel: vi.fn(() => undefined),
+    shouldPreferProviderRuntimeResolvedModel: vi.fn(() => false),
+    shouldDeferProviderSyntheticProfileAuthWithPlugin: vi.fn(() => false),
+  }));
+}
+
+export function installEmbeddedRunnerBackoffE2eMocks(
+  options: EmbeddedRunnerBackoffMockOptions,
+): void {
+  vi.doMock("../../infra/backoff.js", () => ({
+    computeBackoff: (
+      policy: { initialMs: number; maxMs: number; factor: number; jitter: number },
+      attempt: number,
+    ) => options.computeBackoff(policy, attempt),
+    sleepWithAbort: (ms: number, abortSignal?: AbortSignal) =>
+      options.sleepWithAbort(ms, abortSignal),
+  }));
+}