From 3b4f411f7cefc8468ec415a78600a89f22bd143b Mon Sep 17 00:00:00 2001 From: Vincent Koc Date: Sun, 12 Apr 2026 09:49:22 +0100 Subject: [PATCH] test(agents): share mutable auth controller harness --- .../run/auth-controller.test.ts | 221 ++++++++---------- 1 file changed, 98 insertions(+), 123 deletions(-) diff --git a/src/agents/pi-embedded-runner/run/auth-controller.test.ts b/src/agents/pi-embedded-runner/run/auth-controller.test.ts index 02b77f94119..8b54fbcfe39 100644 --- a/src/agents/pi-embedded-runner/run/auth-controller.test.ts +++ b/src/agents/pi-embedded-runner/run/auth-controller.test.ts @@ -62,6 +62,82 @@ function getRuntimeAuthSnapshot( return state ? { profileId: state.profileId, refreshInFlight: state.refreshInFlight } : null; } +type MutableAuthControllerHarness = { + runtimeModel: Model; + effectiveModel: Model; + apiKeyInfo: unknown; + lastProfileId?: string; + runtimeAuthState: RuntimeAuthState | null; + profileIndex: number; +}; + +function createMutableAuthControllerHarness(): MutableAuthControllerHarness { + return { + runtimeModel: createTestModel(), + effectiveModel: createTestModel(), + apiKeyInfo: null, + lastProfileId: undefined, + runtimeAuthState: null, + profileIndex: 0, + }; +} + +function createMutableEmbeddedRunAuthController(params: { + harness: MutableAuthControllerHarness; + setRuntimeApiKey: ReturnType; + profileCandidates?: string[]; +}) { + return createEmbeddedRunAuthController({ + config: undefined, + agentDir: "/tmp/agent", + workspaceDir: "/tmp/workspace", + authStore: { + version: 1, + profiles: {}, + } as AuthProfileStore, + authStorage: { setRuntimeApiKey: params.setRuntimeApiKey }, + profileCandidates: params.profileCandidates ?? ["default"], + initialThinkLevel: "medium", + attemptedThinking: new Set(), + fallbackConfigured: false, + allowTransientCooldownProbe: false, + getProvider: () => "custom-openai", + getModelId: () => "test-model", + getRuntimeModel: () => params.harness.runtimeModel, + setRuntimeModel: (next) => { + params.harness.runtimeModel = next; + }, + getEffectiveModel: () => params.harness.effectiveModel, + setEffectiveModel: (next) => { + params.harness.effectiveModel = next; + }, + getApiKeyInfo: () => params.harness.apiKeyInfo as never, + setApiKeyInfo: (next) => { + params.harness.apiKeyInfo = next; + }, + getLastProfileId: () => params.harness.lastProfileId, + setLastProfileId: (next) => { + params.harness.lastProfileId = next; + }, + getRuntimeAuthState: () => params.harness.runtimeAuthState as never, + setRuntimeAuthState: (next) => { + params.harness.runtimeAuthState = next; + }, + getRuntimeAuthRefreshCancelled: () => false, + setRuntimeAuthRefreshCancelled: () => undefined, + getProfileIndex: () => params.harness.profileIndex, + setProfileIndex: (next) => { + params.harness.profileIndex = next; + }, + setThinkLevel: () => undefined, + log: { + debug: () => undefined, + info: () => undefined, + warn: () => undefined, + }, + }); +} + describe("createEmbeddedRunAuthController", () => { beforeEach(() => { mocks.prepareProviderRuntimeAuth.mockReset(); @@ -69,16 +145,7 @@ describe("createEmbeddedRunAuthController", () => { }); it("applies runtime request overrides on the first auth exchange", async () => { - let runtimeModel = createTestModel(); - let effectiveModel = createTestModel(); - let runtimeAuthState: { - sourceApiKey: string; - authMode: string; - profileId?: string; - expiresAt?: number; - } | null = null; - let apiKeyInfo: unknown = null; - let lastProfileId: string | undefined; + const harness = createMutableAuthControllerHarness(); const setRuntimeApiKey = vi.fn(); mocks.getApiKeyForModel.mockResolvedValue({ @@ -99,66 +166,23 @@ describe("createEmbeddedRunAuthController", () => { }, }); - const controller = createEmbeddedRunAuthController({ - config: undefined, - agentDir: "/tmp/agent", - workspaceDir: "/tmp/workspace", - authStore: { - version: 1, - profiles: {}, - } as AuthProfileStore, - authStorage: { setRuntimeApiKey }, - profileCandidates: ["default"], - initialThinkLevel: "medium", - attemptedThinking: new Set(), - fallbackConfigured: false, - allowTransientCooldownProbe: false, - getProvider: () => "custom-openai", - getModelId: () => "test-model", - getRuntimeModel: () => runtimeModel, - setRuntimeModel: (next) => { - runtimeModel = next; - }, - getEffectiveModel: () => effectiveModel, - setEffectiveModel: (next) => { - effectiveModel = next; - }, - getApiKeyInfo: () => apiKeyInfo as never, - setApiKeyInfo: (next) => { - apiKeyInfo = next; - }, - getLastProfileId: () => lastProfileId, - setLastProfileId: (next) => { - lastProfileId = next; - }, - getRuntimeAuthState: () => runtimeAuthState as never, - setRuntimeAuthState: (next) => { - runtimeAuthState = next; - }, - getRuntimeAuthRefreshCancelled: () => false, - setRuntimeAuthRefreshCancelled: () => undefined, - getProfileIndex: () => 0, - setProfileIndex: () => undefined, - setThinkLevel: () => undefined, - log: { - debug: () => undefined, - info: () => undefined, - warn: () => undefined, - }, + const controller = createMutableEmbeddedRunAuthController({ + harness, + setRuntimeApiKey, }); await controller.initializeAuthProfile(); - expect(runtimeModel.baseUrl).toBe("https://runtime.example.com/v1"); - expect(runtimeModel.headers).toEqual({ + expect(harness.runtimeModel.baseUrl).toBe("https://runtime.example.com/v1"); + expect(harness.runtimeModel.headers).toEqual({ "api-key": "runtime-header-token", }); - expect(effectiveModel.baseUrl).toBe("https://runtime.example.com/v1"); - expect(effectiveModel.headers).toEqual({ + expect(harness.effectiveModel.baseUrl).toBe("https://runtime.example.com/v1"); + expect(harness.effectiveModel.headers).toEqual({ "api-key": "runtime-header-token", }); expect(setRuntimeApiKey).toHaveBeenCalledWith("custom-openai", "runtime-api-key"); - expect(runtimeAuthState).toMatchObject({ + expect(harness.runtimeAuthState).toMatchObject({ sourceApiKey: "source-api-key", authMode: "api-key", profileId: "default", @@ -232,12 +256,7 @@ describe("createEmbeddedRunAuthController", () => { it("ignores stale scheduled refresh results after auth profile rotation", async () => { vi.useFakeTimers(); try { - let runtimeModel = createTestModel(); - let effectiveModel = createTestModel(); - let runtimeAuthState: RuntimeAuthState | null = null; - let apiKeyInfo: unknown = null; - let lastProfileId: string | undefined; - let profileIndex = 0; + const harness = createMutableAuthControllerHarness(); const setRuntimeApiKey = vi.fn(); const staleRefresh = createDeferred<{ apiKey: string; @@ -270,7 +289,7 @@ describe("createEmbeddedRunAuthController", () => { }); mocks.prepareProviderRuntimeAuth.mockImplementation(async ({ context }) => { if (context.apiKey === "default-source-api-key" && context.profileId === "default") { - if (runtimeAuthState?.refreshInFlight) { + if (harness.runtimeAuthState?.refreshInFlight) { return staleRefresh.promise; } return { @@ -303,67 +322,23 @@ describe("createEmbeddedRunAuthController", () => { throw new Error(`Unexpected runtime auth request for ${String(context.profileId)}`); }); - const controller = createEmbeddedRunAuthController({ - config: undefined, - agentDir: "/tmp/agent", - workspaceDir: "/tmp/workspace", - authStore: { - version: 1, - profiles: {}, - } as AuthProfileStore, - authStorage: { setRuntimeApiKey }, + const controller = createMutableEmbeddedRunAuthController({ + harness, + setRuntimeApiKey, profileCandidates: ["default", "backup"], - initialThinkLevel: "medium", - attemptedThinking: new Set(), - fallbackConfigured: false, - allowTransientCooldownProbe: false, - getProvider: () => "custom-openai", - getModelId: () => "test-model", - getRuntimeModel: () => runtimeModel, - setRuntimeModel: (next) => { - runtimeModel = next; - }, - getEffectiveModel: () => effectiveModel, - setEffectiveModel: (next) => { - effectiveModel = next; - }, - getApiKeyInfo: () => apiKeyInfo as never, - setApiKeyInfo: (next) => { - apiKeyInfo = next; - }, - getLastProfileId: () => lastProfileId, - setLastProfileId: (next) => { - lastProfileId = next; - }, - getRuntimeAuthState: () => runtimeAuthState as never, - setRuntimeAuthState: (next) => { - runtimeAuthState = next; - }, - getRuntimeAuthRefreshCancelled: () => false, - setRuntimeAuthRefreshCancelled: () => undefined, - getProfileIndex: () => profileIndex, - setProfileIndex: (next) => { - profileIndex = next; - }, - setThinkLevel: () => undefined, - log: { - debug: () => undefined, - info: () => undefined, - warn: () => undefined, - }, }); await controller.initializeAuthProfile(); - expect(getRuntimeAuthSnapshot(runtimeAuthState)?.profileId).toBe("default"); + expect(getRuntimeAuthSnapshot(harness.runtimeAuthState)?.profileId).toBe("default"); vi.advanceTimersByTime(5_000); await Promise.resolve(); - expect(getRuntimeAuthSnapshot(runtimeAuthState)?.refreshInFlight).toBeTruthy(); + expect(getRuntimeAuthSnapshot(harness.runtimeAuthState)?.refreshInFlight).toBeTruthy(); await controller.advanceAuthProfile(); - expect(getRuntimeAuthSnapshot(runtimeAuthState)?.profileId).toBe("backup"); - expect(runtimeModel.baseUrl).toBe("https://backup-runtime.example.com/v1"); - expect(runtimeModel.headers).toEqual({ + expect(getRuntimeAuthSnapshot(harness.runtimeAuthState)?.profileId).toBe("backup"); + expect(harness.runtimeModel.baseUrl).toBe("https://backup-runtime.example.com/v1"); + expect(harness.runtimeModel.headers).toEqual({ "api-key": "backup-runtime-header-token", }); @@ -382,9 +357,9 @@ describe("createEmbeddedRunAuthController", () => { await Promise.resolve(); await Promise.resolve(); - expect(getRuntimeAuthSnapshot(runtimeAuthState)?.profileId).toBe("backup"); - expect(runtimeModel.baseUrl).toBe("https://backup-runtime.example.com/v1"); - expect(runtimeModel.headers).toEqual({ + expect(getRuntimeAuthSnapshot(harness.runtimeAuthState)?.profileId).toBe("backup"); + expect(harness.runtimeModel.baseUrl).toBe("https://backup-runtime.example.com/v1"); + expect(harness.runtimeModel.headers).toEqual({ "api-key": "backup-runtime-header-token", }); expect(setRuntimeApiKey).toHaveBeenLastCalledWith("custom-openai", "backup-runtime-api-key");