From c9b98875830a147e2819c2ec858b1018214be653 Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Mon, 27 Apr 2026 11:50:33 +0100 Subject: [PATCH] test: speed up embedded runner e2e mocks --- .../model-fallback.run-embedded.e2e.test.ts | 45 ++---- src/agents/pi-embedded-runner.e2e.test.ts | 20 +-- ...pi-agent.auth-profile-rotation.e2e.test.ts | 146 +++--------------- .../pi-embedded-runner-e2e-mocks.ts | 131 ++++++++++++++++ 4 files changed, 174 insertions(+), 168 deletions(-) diff --git a/src/agents/model-fallback.run-embedded.e2e.test.ts b/src/agents/model-fallback.run-embedded.e2e.test.ts index d389d049d50..c88923a2823 100644 --- 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>(); 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( - "./pi-embedded-runner/run/attempt.js", - ); - return { - ...actual, - runEmbeddedAttempt: (params: unknown) => runEmbeddedAttemptMock(params), - }; -}); - -vi.mock("../infra/backoff.js", async () => { - const actual = await vi.importActual("../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("./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( - "../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; diff --git a/src/agents/pi-embedded-runner.e2e.test.ts b/src/agents/pi-embedded-runner.e2e.test.ts index 7d7351e751d..ca14c213f09 100644 --- 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>(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("./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( - "../plugins/provider-runtime.js", - ); - return { - ...actual, - prepareProviderRuntimeAuth: vi.fn(async () => undefined), - }; - }); vi.doMock("./models-config.js", async () => { const mod = await vi.importActual("./models-config.js"); return { diff --git 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 index df152ff473a..f53a0da15d6 100644 --- 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>(); 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"); diff --git a/src/agents/test-helpers/pi-embedded-runner-e2e-mocks.ts b/src/agents/test-helpers/pi-embedded-runner-e2e-mocks.ts index 2e4cee1dadf..93aaaee5b90 100644 --- 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), + })); +}