From 53239102f8e1f2f1f3b29c8d6e14b2ee85fabf6f Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Sat, 18 Apr 2026 17:37:19 +0100 Subject: [PATCH] test: speed up agent model auth tests --- src/agents/model-auth.profiles.test.ts | 82 ++++++++- src/agents/model-fallback.test.ts | 171 +++++++++++++++--- .../models-config.providers.ollama.test.ts | 30 +-- src/agents/pi-model-discovery.auth.test.ts | 17 +- 4 files changed, 256 insertions(+), 44 deletions(-) diff --git a/src/agents/model-auth.profiles.test.ts b/src/agents/model-auth.profiles.test.ts index d04f8c64ffa..d1e371815c9 100644 --- a/src/agents/model-auth.profiles.test.ts +++ b/src/agents/model-auth.profiles.test.ts @@ -5,7 +5,10 @@ import type { Api, Model } from "@mariozechner/pi-ai"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import type { OpenClawConfig } from "../config/types.openclaw.js"; import { withEnvAsync } from "../test-utils/env.js"; -import { clearRuntimeAuthProfileStoreSnapshots, ensureAuthProfileStore } from "./auth-profiles.js"; +import { + clearRuntimeAuthProfileStoreSnapshots, + ensureAuthProfileStore, +} from "./auth-profiles/store.js"; import { getApiKeyForModel, hasAvailableAuthForProvider, @@ -13,6 +16,78 @@ import { resolveEnvApiKey, } from "./model-auth.js"; +vi.mock("../plugins/setup-registry.js", async () => { + const { readFileSync } = await import("node:fs"); + return { + resolvePluginSetupProvider: ({ provider }: { provider: string; env: NodeJS.ProcessEnv }) => { + if (provider !== "anthropic-vertex") { + return undefined; + } + return { + resolveConfigApiKey: ({ env }: { env: NodeJS.ProcessEnv }) => { + const metadataOptIn = env.ANTHROPIC_VERTEX_USE_GCP_METADATA?.trim().toLowerCase(); + if (metadataOptIn === "1" || metadataOptIn === "true") { + return "gcp-vertex-credentials"; + } + const credentialsPath = env.GOOGLE_APPLICATION_CREDENTIALS?.trim(); + if (!credentialsPath) { + return undefined; + } + try { + readFileSync(credentialsPath, "utf8"); + return "gcp-vertex-credentials"; + } catch { + return undefined; + } + }, + }; + }, + }; +}); + +vi.mock("./provider-auth-aliases.js", () => ({ + resolveProviderAuthAliasMap: () => ({}), + resolveProviderIdForAuth: (provider: string) => { + const normalized = provider.trim().toLowerCase(); + if (normalized === "modelstudio" || normalized === "qwencloud") { + return "qwen"; + } + if (normalized === "z.ai" || normalized === "z-ai") { + return "zai"; + } + if (normalized === "opencode-go-auth") { + return "opencode-go"; + } + if (normalized === "bedrock" || normalized === "aws-bedrock") { + return "amazon-bedrock"; + } + return normalized; + }, +})); + +vi.mock("./model-auth-env-vars.js", () => { + const candidates = { + anthropic: ["ANTHROPIC_OAUTH_TOKEN", "ANTHROPIC_API_KEY"], + google: ["GEMINI_API_KEY", "GOOGLE_API_KEY"], + huggingface: ["HUGGINGFACE_HUB_TOKEN", "HF_TOKEN"], + "minimax-portal": ["MINIMAX_OAUTH_TOKEN", "MINIMAX_API_KEY"], + ollama: ["OLLAMA_API_KEY"], + "opencode-go": ["OPENCODE_API_KEY", "OPENCODE_ZEN_API_KEY"], + openai: ["OPENAI_API_KEY"], + qianfan: ["QIANFAN_API_KEY"], + qwen: ["QWEN_API_KEY", "MODELSTUDIO_API_KEY", "DASHSCOPE_API_KEY"], + synthetic: ["SYNTHETIC_API_KEY"], + "vercel-ai-gateway": ["AI_GATEWAY_API_KEY"], + voyage: ["VOYAGE_API_KEY"], + zai: ["ZAI_API_KEY", "Z_AI_API_KEY"], + } as const; + return { + PROVIDER_ENV_API_KEY_CANDIDATES: candidates, + listKnownProviderEnvApiKeyNames: () => [...new Set(Object.values(candidates).flat())], + resolveProviderEnvApiKeyCandidates: () => candidates, + }; +}); + vi.mock("../plugins/provider-runtime.js", () => ({ buildProviderMissingAuthMessageWithPlugin: (params: { provider: string; @@ -68,6 +143,11 @@ vi.mock("../plugins/provider-runtime.js", () => ({ }, })); +vi.mock("../plugins/providers.js", () => ({ + resolveOwningPluginIdsForProvider: ({ provider }: { provider: string }) => + provider === "openai" ? ["openai"] : [], +})); + vi.mock("./cli-credentials.js", () => ({ readCodexCliCredentialsCached: () => null, readMiniMaxCliCredentialsCached: () => null, diff --git a/src/agents/model-fallback.test.ts b/src/agents/model-fallback.test.ts index beb86bf8f40..b128224fe06 100644 --- a/src/agents/model-fallback.test.ts +++ b/src/agents/model-fallback.test.ts @@ -2,14 +2,17 @@ import crypto from "node:crypto"; import fs from "node:fs/promises"; import os from "node:os"; import path from "node:path"; -import { afterAll, beforeAll, describe, expect, it, vi } from "vitest"; +import { afterAll, afterEach, beforeAll, describe, expect, it, vi } from "vitest"; import type { OpenClawConfig } from "../config/config.js"; import { resetLogger, setLoggerOverride } from "../logging/logger.js"; import { createWarnLogCapture } from "../logging/test-helpers/warn-log-capture.js"; import { AUTH_STORE_VERSION } from "./auth-profiles/constants.js"; import * as authProfileSourceCheckModule from "./auth-profiles/source-check.js"; import * as authProfileStoreModule from "./auth-profiles/store.js"; -import { saveAuthProfileStore } from "./auth-profiles/store.js"; +import { + clearRuntimeAuthProfileStoreSnapshots, + replaceRuntimeAuthProfileStoreSnapshots, +} from "./auth-profiles/store.js"; import type { AuthProfileStore } from "./auth-profiles/types.js"; import { isAnthropicBillingError } from "./live-auth-keys.js"; import { LiveSessionModelSwitchError } from "./live-model-switch-error.js"; @@ -25,6 +28,111 @@ vi.mock("../plugins/provider-runtime.js", () => ({ resolveExternalAuthProfilesWithPlugins: () => [], })); +const authRuntimeMock = vi.hoisted(() => { + const stores = new Map(); + const keyFor = (agentDir?: string) => agentDir ?? "__main__"; + const now = () => Date.now(); + const isActive = (value: unknown, ts = now()) => + typeof value === "number" && Number.isFinite(value) && value > ts; + const cloneStore = (store: AuthProfileStore): AuthProfileStore => structuredClone(store); + const getStore = (agentDir?: string): AuthProfileStore => + cloneStore(stores.get(keyFor(agentDir)) ?? { version: 1, profiles: {} }); + const getProfileIds = (store: AuthProfileStore, provider: string) => + Object.entries(store.profiles) + .filter(([, profile]) => profile.provider === provider) + .map(([id]) => id); + const isProfileInCooldown = ( + store: AuthProfileStore, + profileId: string, + tsOrOptions?: number | { now?: number; forModel?: string }, + forModel?: string, + ) => { + const stats = store.usageStats?.[profileId]; + if (!stats || store.profiles[profileId]?.provider === "openrouter") { + return false; + } + const ts = typeof tsOrOptions === "number" ? tsOrOptions : (tsOrOptions?.now ?? now()); + const model = typeof tsOrOptions === "object" ? tsOrOptions.forModel : forModel; + if (isActive(stats.disabledUntil, ts)) { + return true; + } + if (!isActive(stats.cooldownUntil, ts)) { + return false; + } + return !stats.cooldownModel || !model || stats.cooldownModel === model; + }; + const resolveReason = (store: AuthProfileStore, profileIds: string[], ts = now()) => { + for (const profileId of profileIds) { + const stats = store.usageStats?.[profileId]; + if (!stats) { + continue; + } + if (isActive(stats.disabledUntil, ts)) { + return stats.disabledReason ?? "auth"; + } + if (!isActive(stats.cooldownUntil, ts)) { + continue; + } + if (stats.cooldownReason) { + return stats.cooldownReason; + } + const counts = stats.failureCounts ?? {}; + if ((counts.rate_limit ?? 0) > 0) { + return "rate_limit"; + } + if ((counts.overloaded ?? 0) > 0) { + return "overloaded"; + } + if ((counts.timeout ?? 0) > 0) { + return "timeout"; + } + return "unknown"; + } + return null; + }; + return { + clear: () => stores.clear(), + setStore: (agentDir: string | undefined, store: AuthProfileStore) => { + stores.set(keyFor(agentDir), cloneStore(store)); + }, + runtime: { + ensureAuthProfileStore: (agentDir?: string) => getStore(agentDir), + loadAuthProfileStoreForRuntime: (agentDir?: string) => getStore(agentDir), + resolveAuthProfileOrder: (params: { store: AuthProfileStore; provider: string }) => + getProfileIds(params.store, params.provider), + isProfileInCooldown, + resolveProfilesUnavailableReason: (params: { + store: AuthProfileStore; + profileIds: string[]; + now?: number; + }) => resolveReason(params.store, params.profileIds, params.now), + getSoonestCooldownExpiry: ( + store: AuthProfileStore, + profileIds: string[], + options?: { now?: number; forModel?: string }, + ) => { + const ts = options?.now ?? now(); + let soonest: number | null = null; + for (const profileId of profileIds) { + if (!isProfileInCooldown(store, profileId, { now: ts, forModel: options?.forModel })) { + continue; + } + const stats = store.usageStats?.[profileId]; + const expiry = [stats?.cooldownUntil, stats?.disabledUntil] + .filter((value): value is number => isActive(value, ts)) + .toSorted((a, b) => a - b)[0]; + if (expiry !== undefined && (soonest === null || expiry < soonest)) { + soonest = expiry; + } + } + return soonest; + }, + }, + }; +}); + +vi.mock("./model-fallback-auth.runtime.js", () => authRuntimeMock.runtime); + const makeCfg = makeModelFallbackCfg; const OPENROUTER_MODEL_NOT_FOUND_PAYLOAD = '{"error":{"message":"Healer Alpha was a stealth model revealed on March 18th as an early testing version of MiMo-V2-Omni. Find it here: https://openrouter.ai/xiaomi/mimo-v2-omni","code":404},"user_id":"user_33GTyP8uDSYYbaeBO48AGHXyuMC"}'; @@ -41,6 +149,11 @@ afterAll(async () => { } }); +afterEach(() => { + clearRuntimeAuthProfileStoreSnapshots(); + authRuntimeMock.clear(); +}); + function makeFallbacksOnlyCfg(): OpenClawConfig { return { agents: { @@ -71,7 +184,7 @@ async function withTempAuthStore( run: (tempDir: string) => Promise, ): Promise { const tempDir = await makeAuthTempDir(); - saveAuthProfileStore(store, tempDir); + setAuthRuntimeStore(tempDir, store); return await run(tempDir); } @@ -87,15 +200,20 @@ async function runWithStoredAuth(params: { provider: string; run: (provider: string, model: string) => Promise; }) { - return withTempAuthStore(params.store, async (tempDir) => - runWithModelFallback({ - cfg: params.cfg, - provider: params.provider, - model: "m1", - agentDir: tempDir, - run: params.run, - }), - ); + const tempDir = await makeAuthTempDir(); + setAuthRuntimeStore(tempDir, params.store); + return await runWithModelFallback({ + cfg: params.cfg, + provider: params.provider, + model: "m1", + agentDir: tempDir, + run: params.run, + }); +} + +function setAuthRuntimeStore(agentDir: string | undefined, store: AuthProfileStore): void { + replaceRuntimeAuthProfileStoreSnapshots([{ agentDir, store }]); + authRuntimeMock.setStore(agentDir, store); } async function expectFallsBackToHaiku(params: { @@ -853,20 +971,17 @@ describe("runWithModelFallback", () => { await withTempAuthStore(store, async (tempDir) => { const run = vi.fn().mockImplementation(async (provider: string, model: string) => { if (provider === "anthropic" && model === "claude-opus-4-5") { - saveAuthProfileStore( - { - ...store, - usageStats: { - "anthropic:default": { - cooldownUntil: expiry, - cooldownReason: "rate_limit", - cooldownModel: "claude-opus-4-5", - failureCounts: { rate_limit: 1 }, - }, + setAuthRuntimeStore(tempDir, { + ...store, + usageStats: { + "anthropic:default": { + cooldownUntil: expiry, + cooldownReason: "rate_limit", + cooldownModel: "claude-opus-4-5", + failureCounts: { rate_limit: 1 }, }, }, - tempDir, - ); + }); } throw Object.assign(new Error("rate limited"), { status: 429 }); @@ -1386,7 +1501,7 @@ describe("runWithModelFallback", () => { async function makeAuthStoreWithCooldown( provider: string, reason: "rate_limit" | "overloaded" | "timeout" | "auth" | "billing", - ): Promise<{ store: AuthProfileStore; dir: string }> { + ): Promise<{ dir: string }> { const tmpDir = await makeAuthTempDir(); const now = Date.now(); const store: AuthProfileStore = { @@ -1407,8 +1522,8 @@ describe("runWithModelFallback", () => { }, }, }; - saveAuthProfileStore(store, tmpDir); - return { store, dir: tmpDir }; + setAuthRuntimeStore(tmpDir, store); + return { dir: tmpDir }; } it("attempts same-provider fallbacks during rate limit cooldown", async () => { @@ -1572,7 +1687,7 @@ describe("runWithModelFallback", () => { }, }, }; - saveAuthProfileStore(store, tmpDir); + setAuthRuntimeStore(tmpDir, store); const cfg = makeCfg({ agents: { diff --git a/src/agents/models-config.providers.ollama.test.ts b/src/agents/models-config.providers.ollama.test.ts index ddd074ab32e..13ec724069a 100644 --- a/src/agents/models-config.providers.ollama.test.ts +++ b/src/agents/models-config.providers.ollama.test.ts @@ -9,11 +9,17 @@ import { runProviderCatalog, } from "../plugins/provider-discovery.js"; import type { ProviderPlugin } from "../plugins/types.js"; -import { loadBundledPluginPublicSurfaceSync } from "../test-utils/bundled-plugin-public-surface.js"; +import { resolveRelativeBundledPluginPublicModuleId } from "../test-utils/bundled-plugin-public-surface.js"; import { withFetchPreconnect } from "../test-utils/fetch-mock.js"; import { OLLAMA_LOCAL_AUTH_MARKER } from "./model-auth-markers.js"; import type { ProviderConfig } from "./models-config.providers.secrets.js"; +const OLLAMA_PROVIDER_DISCOVERY_MODULE_ID = resolveRelativeBundledPluginPublicModuleId({ + fromModuleUrl: import.meta.url, + pluginId: "ollama", + artifactBasename: "provider-discovery.js", +}); + afterEach(() => { vi.unstubAllEnvs(); vi.unstubAllGlobals(); @@ -48,20 +54,16 @@ describe("Ollama provider", () => { } } - let ollamaCatalogProvider: ProviderPlugin | undefined; + let ollamaCatalogProvider: Promise | undefined; - function loadOllamaCatalogProvider(): ProviderPlugin | undefined { - if (ollamaCatalogProvider) { - return ollamaCatalogProvider; - } - const surface = loadBundledPluginPublicSurfaceSync<{ - default?: ProviderPlugin; - ollamaProviderDiscovery?: ProviderPlugin; - }>({ - pluginId: "ollama", - artifactBasename: "provider-discovery.js", + function loadOllamaCatalogProvider(): Promise { + ollamaCatalogProvider ??= import(OLLAMA_PROVIDER_DISCOVERY_MODULE_ID).then((surface) => { + const typed = surface as { + default?: ProviderPlugin; + ollamaProviderDiscovery?: ProviderPlugin; + }; + return typed.default ?? typed.ollamaProviderDiscovery; }); - ollamaCatalogProvider = surface.default ?? surface.ollamaProviderDiscovery; return ollamaCatalogProvider; } @@ -69,7 +71,7 @@ describe("Ollama provider", () => { config?: OpenClawConfig; env?: NodeJS.ProcessEnv; }): Promise { - const provider = loadOllamaCatalogProvider(); + const provider = await loadOllamaCatalogProvider(); if (!provider) { return undefined; } diff --git a/src/agents/pi-model-discovery.auth.test.ts b/src/agents/pi-model-discovery.auth.test.ts index 65ab37d3e6a..dd187ffff61 100644 --- a/src/agents/pi-model-discovery.auth.test.ts +++ b/src/agents/pi-model-discovery.auth.test.ts @@ -1,13 +1,28 @@ import fs from "node:fs/promises"; import os from "node:os"; import path from "node:path"; -import { describe, expect, it } from "vitest"; +import { describe, expect, it, vi } from "vitest"; import { resolvePiCredentialMapFromStore } from "./pi-auth-credentials.js"; import { addEnvBackedPiCredentials, scrubLegacyStaticAuthJsonEntriesForDiscovery, } from "./pi-model-discovery.js"; +vi.mock("./model-auth-env-vars.js", () => ({ + resolveProviderEnvApiKeyCandidates: () => ({ + mistral: ["MISTRAL_API_KEY"], + }), +})); + +vi.mock("./model-auth-env.js", () => ({ + resolveEnvApiKey: (provider: string, env: NodeJS.ProcessEnv) => { + if (provider !== "mistral" || !env.MISTRAL_API_KEY?.trim()) { + return null; + } + return { apiKey: env.MISTRAL_API_KEY, source: "env: MISTRAL_API_KEY" }; + }, +})); + async function createAgentDir(): Promise { return await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-pi-auth-storage-")); }