From c1fc2ed0e882dd3255ca119eaddc8c0144407d03 Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Tue, 7 Apr 2026 17:36:09 +0100 Subject: [PATCH] test: speed up provider auth onboarding test --- ...oard-non-interactive.provider-auth.test.ts | 241 ++++++++++++++---- 1 file changed, 186 insertions(+), 55 deletions(-) diff --git a/src/commands/onboard-non-interactive.provider-auth.test.ts b/src/commands/onboard-non-interactive.provider-auth.test.ts index 16ba360cdba..5cbfcc5855e 100644 --- a/src/commands/onboard-non-interactive.provider-auth.test.ts +++ b/src/commands/onboard-non-interactive.provider-auth.test.ts @@ -22,26 +22,83 @@ const OPENAI_DEFAULT_MODEL = "openai/gpt-5.4"; const ZAI_CODING_GLOBAL_BASE_URL = "https://api.z.ai/api/coding/paas/v4"; const ZAI_CODING_CN_BASE_URL = "https://open.bigmodel.cn/api/coding/paas/v4"; const ZAI_GLOBAL_BASE_URL = "https://api.z.ai/api/paas/v4"; +const TEST_AUTH_STORE_VERSION = 1; +const TEST_MAIN_AUTH_STORE_KEY = "__main__"; const ensureWorkspaceAndSessionsMock = vi.hoisted(() => vi.fn(async (..._args: unknown[]) => {})); +const testAuthProfileStores = vi.hoisted( + () => new Map> }>(), +); + +function normalizeStoredSecret(value: unknown): string { + return typeof value === "string" ? value.replaceAll("\r", "").replaceAll("\n", "").trim() : ""; +} + +function cloneTestAuthStore(store: { + version: number; + profiles: Record>; +}) { + return structuredClone(store); +} + +function writeRuntimeAuthSnapshots() { + if (!replaceRuntimeAuthProfileStoreSnapshots) { + return; + } + replaceRuntimeAuthProfileStoreSnapshots( + Array.from(testAuthProfileStores.entries()).map(([key, store]) => + key === TEST_MAIN_AUTH_STORE_KEY + ? { store: cloneTestAuthStore(store) as never } + : { agentDir: key, store: cloneTestAuthStore(store) as never }, + ), + ); +} + +function getOrCreateTestAuthStore(agentDir?: string) { + const key = agentDir?.trim() || TEST_MAIN_AUTH_STORE_KEY; + let store = testAuthProfileStores.get(key); + if (!store) { + store = { version: TEST_AUTH_STORE_VERSION, profiles: {} }; + testAuthProfileStores.set(key, store); + } + return store; +} + +function upsertAuthProfile(params: { + profileId: string; + credential: Record; + agentDir?: string; +}) { + const credential = + params.credential.type === "api_key" && typeof params.credential.key === "string" + ? { + ...params.credential, + key: normalizeStoredSecret(params.credential.key), + } + : params.credential.type === "token" && typeof params.credential.token === "string" + ? { + ...params.credential, + token: normalizeStoredSecret(params.credential.token), + } + : params.credential; + for (const targetAgentDir of new Set([undefined, params.agentDir])) { + const store = getOrCreateTestAuthStore(targetAgentDir); + store.profiles[params.profileId] = credential; + } + writeRuntimeAuthSnapshots(); +} vi.mock("./onboard-non-interactive/local/auth-choice.plugin-providers.js", async () => { const [ { resolveDefaultAgentId, resolveAgentDir, resolveAgentWorkspaceDir }, { resolveDefaultAgentWorkspaceDir }, { enablePluginInConfig }, - { upsertAuthProfile }, - { createProviderApiKeyAuthMethod }, - { providerApiKeyAuthRuntime }, { configureOpenAICompatibleSelfHostedProviderNonInteractive }, { detectZaiEndpoint }, ] = await Promise.all([ import("../agents/agent-scope.js"), import("../agents/workspace.js"), import("../plugins/enable.js"), - import("../agents/auth-profiles/profiles.js"), - import("../plugins/provider-api-key-auth.js"), - import("../plugins/provider-api-key-auth.runtime.js"), import("../plugins/provider-self-hosted-setup.js"), import("../plugins/provider-zai-endpoint.js"), ]); @@ -164,6 +221,68 @@ vi.mock("./onboard-non-interactive/local/auth-choice.plugin-providers.js", async }; } + function applyAuthProfileConfig( + cfg: Record, + params: { + profileId: string; + provider: string; + mode: "api_key" | "oauth" | "token"; + email?: string; + displayName?: string; + }, + ): Record { + const auth = + cfg.auth && typeof cfg.auth === "object" ? (cfg.auth as Record) : {}; + const profiles = + auth.profiles && typeof auth.profiles === "object" + ? (auth.profiles as Record) + : {}; + return { + ...cfg, + auth: { + ...auth, + profiles: { + ...profiles, + [params.profileId]: { + provider: params.provider, + mode: params.mode, + ...(params.email ? { email: params.email } : {}), + ...(params.displayName ? { displayName: params.displayName } : {}), + }, + }, + }, + }; + } + + function applyPrimaryModel(cfg: Record, model: string): Record { + const agents = + cfg.agents && typeof cfg.agents === "object" ? (cfg.agents as Record) : {}; + const defaults = + agents.defaults && typeof agents.defaults === "object" + ? (agents.defaults as Record) + : {}; + const models = + defaults.models && typeof defaults.models === "object" + ? (defaults.models as Record) + : {}; + return { + ...cfg, + agents: { + ...agents, + defaults: { + ...defaults, + model: { + primary: model, + }, + models: { + ...models, + [model]: models[model] ?? {}, + }, + }, + }, + }; + } + function createApiKeyChoice(params: { providerId: string; label: string; @@ -177,32 +296,52 @@ vi.mock("./onboard-non-interactive/local/auth-choice.plugin-providers.js", async profileIds?: string[]; applyConfig?: (cfg: Record) => Record; }): ChoiceHandler { - const method = createProviderApiKeyAuthMethod({ - providerId: params.providerId, - methodId: "api-key", - label: params.label, - optionKey: params.optionKey, - flagName: params.flagName, - envVar: params.envVar, - promptMessage: `Enter ${params.label} API key`, - ...(params.profileId ? { profileId: params.profileId } : {}), - ...(params.profileIds ? { profileIds: params.profileIds } : {}), - ...(params.defaultModel ? { defaultModel: params.defaultModel } : {}), - ...(params.applyConfig - ? { applyConfig: (cfg) => params.applyConfig!(cfg as Record) as never } - : {}), - wizard: { - choiceId: params.choiceId, - choiceLabel: params.label, - groupId: params.providerId, - groupLabel: params.label, - }, - }); + const profileIds = + params.profileIds?.map((value) => value.trim()).filter(Boolean) ?? + (params.profileId ? [params.profileId] : [`${params.providerId}:default`]); return { providerId: params.providerId, label: params.label, ...(params.pluginId ? { pluginId: params.pluginId } : {}), - runNonInteractive: async (ctx) => await method.runNonInteractive!(ctx as never), + runNonInteractive: async (ctx) => { + const resolved = await ctx.resolveApiKey({ + provider: params.providerId, + flagValue: normalizeText(ctx.opts[params.optionKey]), + flagName: params.flagName, + envVar: params.envVar, + }); + if (!resolved) { + return null; + } + if (resolved.source !== "profile") { + for (const profileId of profileIds) { + const credential = ctx.toApiKeyCredential({ + provider: profileId.split(":", 1)[0]?.trim() || params.providerId, + resolved, + }); + if (!credential) { + return null; + } + upsertAuthProfile({ + profileId, + credential, + agentDir: ctx.agentDir, + }); + } + } + let next = ctx.config; + for (const profileId of profileIds) { + next = applyAuthProfileConfig(next, { + profileId, + provider: profileId.split(":", 1)[0]?.trim() || params.providerId, + mode: "api_key", + }); + } + if (params.applyConfig) { + next = params.applyConfig(next); + } + return params.defaultModel ? applyPrimaryModel(next, params.defaultModel) : next; + }, }; } @@ -264,11 +403,11 @@ vi.mock("./onboard-non-interactive/local/auth-choice.plugin-providers.js", async : {}), }); const fallback = ZAI_FALLBACKS[choiceId]; - let next = providerApiKeyAuthRuntime.applyAuthProfileConfig(ctx.config as never, { + let next = applyAuthProfileConfig(ctx.config as never, { profileId: "zai:default", provider: "zai", mode: "api_key", - }) as Record; + }); next = withProviderConfig(next, "zai", { baseUrl: detected?.baseUrl ?? fallback.baseUrl, api: "openai-completions", @@ -278,10 +417,7 @@ vi.mock("./onboard-non-interactive/local/auth-choice.plugin-providers.js", async }), ], }); - return providerApiKeyAuthRuntime.applyPrimaryModel( - next as never, - `zai/${detected?.modelId ?? fallback.modelId}`, - ); + return applyPrimaryModel(next as never, `zai/${detected?.modelId ?? fallback.modelId}`); }, }; } @@ -316,15 +452,12 @@ vi.mock("./onboard-non-interactive/local/auth-choice.plugin-providers.js", async agentDir: ctx.agentDir, }); } - const withProfile = providerApiKeyAuthRuntime.applyAuthProfileConfig(ctx.config as never, { + const withProfile = applyAuthProfileConfig(ctx.config as never, { profileId: "cloudflare-ai-gateway:default", provider: "cloudflare-ai-gateway", mode: "api_key", }); - return providerApiKeyAuthRuntime.applyPrimaryModel( - withProfile, - "cloudflare-ai-gateway/claude-sonnet-4-5", - ); + return applyPrimaryModel(withProfile, "cloudflare-ai-gateway/claude-sonnet-4-5"); }, }; @@ -350,18 +483,12 @@ vi.mock("./onboard-non-interactive/local/auth-choice.plugin-providers.js", async } as never, agentDir: ctx.agentDir, }); - const withProfile = providerApiKeyAuthRuntime.applyAuthProfileConfig( - ctx.config as never, - { - profileId: (ctx.opts.tokenProfileId as string | undefined) ?? "anthropic:default", - provider: "anthropic", - mode: "token", - }, - ); - return providerApiKeyAuthRuntime.applyPrimaryModel( - withProfile, - "anthropic/claude-sonnet-4-6", - ); + const withProfile = applyAuthProfileConfig(ctx.config as never, { + profileId: (ctx.opts.tokenProfileId as string | undefined) ?? "anthropic:default", + provider: "anthropic", + mode: "token", + }); + return applyPrimaryModel(withProfile, "anthropic/claude-sonnet-4-6"); }, }, ], @@ -656,7 +783,7 @@ const NON_INTERACTIVE_DEFAULT_OPTIONS = { let runNonInteractiveSetup: typeof import("./onboard-non-interactive.js").runNonInteractiveSetup; let clearRuntimeAuthProfileStoreSnapshots: typeof import("../agents/auth-profiles.js").clearRuntimeAuthProfileStoreSnapshots; let ensureAuthProfileStore: typeof import("../agents/auth-profiles.js").ensureAuthProfileStore; -let upsertAuthProfile: typeof import("../agents/auth-profiles.js").upsertAuthProfile; +let replaceRuntimeAuthProfileStoreSnapshots: typeof import("../agents/auth-profiles.js").replaceRuntimeAuthProfileStoreSnapshots; let resetFileLockStateForTest: typeof import("../infra/file-lock.js").resetFileLockStateForTest; let clearPluginDiscoveryCache: typeof import("../plugins/discovery.js").clearPluginDiscoveryCache; let clearPluginManifestRegistryCache: typeof import("../plugins/manifest-registry.js").clearPluginManifestRegistryCache; @@ -866,10 +993,12 @@ async function expectApiKeyProfile(params: { } async function loadProviderAuthOnboardModules(): Promise { - vi.resetModules(); ({ runNonInteractiveSetup } = await import("./onboard-non-interactive.js")); - ({ clearRuntimeAuthProfileStoreSnapshots, ensureAuthProfileStore, upsertAuthProfile } = - await import("../agents/auth-profiles.js")); + ({ + clearRuntimeAuthProfileStoreSnapshots, + ensureAuthProfileStore, + replaceRuntimeAuthProfileStoreSnapshots, + } = await import("../agents/auth-profiles.js")); ({ resetFileLockStateForTest } = await import("../infra/file-lock.js")); ({ clearPluginDiscoveryCache } = await import("../plugins/discovery.js")); ({ clearPluginManifestRegistryCache } = await import("../plugins/manifest-registry.js")); @@ -881,6 +1010,7 @@ describe("onboard (non-interactive): provider auth", () => { }); beforeEach(() => { + testAuthProfileStores.clear(); clearRuntimeAuthProfileStoreSnapshots(); resetFileLockStateForTest(); clearPluginDiscoveryCache(); @@ -890,6 +1020,7 @@ describe("onboard (non-interactive): provider auth", () => { afterEach(() => { vi.unstubAllGlobals(); + testAuthProfileStores.clear(); clearRuntimeAuthProfileStoreSnapshots(); resetFileLockStateForTest(); clearPluginDiscoveryCache();