test: speed up agent runtime helper tests

src/agents/context.lookup.test.ts

@@ -1,29 +1,46 @@
-import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import { afterEach, beforeAll, beforeEach, describe, expect, it, vi } from "vitest";
 
 type DiscoveredModel = { id: string; contextWindow?: number; contextTokens?: number };
 type ContextModule = typeof import("./context.js");
 
+const contextTestState = vi.hoisted(() => {
+  const state = {
+    loadConfigImpl: () => ({}) as unknown,
+    discoveredModels: [] as DiscoveredModel[],
+    ensureOpenClawModelsJson: vi.fn(async () => {}),
+    discoverAuthStorage: vi.fn(() => ({})),
+    discoverModels: vi.fn(() => ({
+      getAll: () => state.discoveredModels,
+    })),
+  };
+  return state;
+});
+
+vi.mock("../config/config.js", () => ({
+  loadConfig: () => contextTestState.loadConfigImpl(),
+}));
+
+vi.mock("./models-config.js", () => ({
+  ensureOpenClawModelsJson: contextTestState.ensureOpenClawModelsJson,
+}));
+
+vi.mock("./agent-paths.js", () => ({
+  resolveOpenClawAgentDir: () => "/tmp/openclaw-agent",
+}));
+
+vi.mock("./pi-model-discovery-runtime.js", () => ({
+  discoverAuthStorage: contextTestState.discoverAuthStorage,
+  discoverModels: contextTestState.discoverModels,
+}));
+
 function mockContextDeps(params: {
   loadConfig: () => unknown;
   discoveredModels?: DiscoveredModel[];
 }) {
-  const ensureOpenClawModelsJson = vi.fn(async () => {});
-  vi.doMock("../config/config.js", () => ({
-    loadConfig: params.loadConfig,
-  }));
-  vi.doMock("./models-config.js", () => ({
-    ensureOpenClawModelsJson,
-  }));
-  vi.doMock("./agent-paths.js", () => ({
-    resolveOpenClawAgentDir: () => "/tmp/openclaw-agent",
-  }));
-  vi.doMock("./pi-model-discovery-runtime.js", () => ({
-    discoverAuthStorage: vi.fn(() => ({})),
-    discoverModels: vi.fn(() => ({
-      getAll: () => params.discoveredModels ?? [],
-    })),
-  }));
-  return { ensureOpenClawModelsJson };
+  contextTestState.loadConfigImpl = params.loadConfig;
+  contextTestState.discoveredModels = params.discoveredModels ?? [];
+  contextTestState.ensureOpenClawModelsJson.mockClear();
+  return { ensureOpenClawModelsJson: contextTestState.ensureOpenClawModelsJson };
 }
 
 function mockContextModuleDeps(loadConfigImpl: () => unknown) {
@@ -61,11 +78,16 @@ async function flushAsyncWarmup() {
   await new Promise((r) => setTimeout(r, 0));
 }
 
-let lastContextModule: ContextModule | null = null;
+let contextModule: ContextModule;
 
 async function importContextModule(): Promise<ContextModule> {
+  await flushAsyncWarmup();
+  return contextModule;
+}
+
+async function importFreshContextModule(): Promise<ContextModule> {
+  vi.resetModules();
   const module = await import("./context.js");
-  lastContextModule = module;
   await flushAsyncWarmup();
   return module;
 }
@@ -76,15 +98,21 @@ async function importResolveContextTokensForModel() {
 }
 
 describe("lookupContextTokens", () => {
+  beforeAll(async () => {
+    contextModule = await import("./context.js");
+  });
+
   beforeEach(() => {
-    vi.resetModules();
-    lastContextModule = null;
+    contextTestState.loadConfigImpl = () => ({});
+    contextTestState.discoveredModels = [];
+    contextTestState.ensureOpenClawModelsJson.mockClear();
+    contextTestState.discoverAuthStorage.mockClear();
+    contextTestState.discoverModels.mockClear();
+    contextModule.resetContextWindowCacheForTest();
   });
 
   afterEach(async () => {
-    if (lastContextModule) {
-      lastContextModule.resetContextWindowCacheForTest();
-    }
+    contextModule.resetContextWindowCacheForTest();
     await flushAsyncWarmup();
   });
 
@@ -147,7 +175,7 @@ describe("lookupContextTokens", () => {
     }));
     mockContextModuleDeps(firstLoadConfigMock);
 
-    let { lookupContextTokens } = await importContextModule();
+    let { lookupContextTokens } = await importFreshContextModule();
     expect(lookupContextTokens("openrouter/claude-sonnet", { allowAsyncLoad: false })).toBe(
       321_000,
     );
@@ -160,7 +188,7 @@ describe("lookupContextTokens", () => {
     });
     mockContextModuleDeps(secondLoadConfigMock);
 
-    ({ lookupContextTokens } = await importContextModule());
+    ({ lookupContextTokens } = await importFreshContextModule());
     expect(lookupContextTokens("openrouter/claude-sonnet", { allowAsyncLoad: false })).toBe(
       321_000,
     );
@@ -192,11 +220,10 @@ describe("lookupContextTokens", () => {
           expectedCalls: 0,
         },
       ]) {
-        vi.resetModules();
         const loadConfigMock = vi.fn(() => ({ models: {} }));
         const { ensureOpenClawModelsJson } = mockContextModuleDeps(loadConfigMock);
         process.argv = scenario.argv;
-        await importContextModule();
+        await importFreshContextModule();
         expect(loadConfigMock).toHaveBeenCalledTimes(scenario.expectedCalls);
         expect(ensureOpenClawModelsJson).toHaveBeenCalledTimes(scenario.expectedCalls);
       }

src/agents/pi-tool-definition-adapter.after-tool-call.fires-once.test.ts

@@ -8,7 +8,7 @@
  */
 import type { AgentTool } from "@mariozechner/pi-agent-core";
 import { Type } from "@sinclair/typebox";
-import { beforeEach, describe, expect, it, vi } from "vitest";
+import { beforeAll, beforeEach, describe, expect, it, vi } from "vitest";
 import { createBaseToolHandlerState } from "./pi-tool-handler-state.test-helpers.js";
 
 const hookMocks = vi.hoisted(() => ({
@@ -79,7 +79,6 @@ let handleToolExecutionStart: typeof import("./pi-embedded-subscribe.handlers.to
 let handleToolExecutionEnd: typeof import("./pi-embedded-subscribe.handlers.tools.js").handleToolExecutionEnd;
 
 async function loadFreshAfterToolCallModulesForTest() {
-  vi.resetModules();
   vi.doMock("../plugins/hook-runner-global.js", () => ({
     getGlobalHookRunner: () => hookMocks.runner,
   }));
@@ -99,7 +98,9 @@ async function loadFreshAfterToolCallModulesForTest() {
 }
 
 describe("after_tool_call fires exactly once in embedded runs", () => {
-  beforeEach(async () => {
+  beforeAll(loadFreshAfterToolCallModulesForTest);
+
+  beforeEach(() => {
     hookMocks.runner.hasHooks.mockClear();
     hookMocks.runner.hasHooks.mockReturnValue(true);
     hookMocks.runner.runAfterToolCall.mockClear();
@@ -115,7 +116,6 @@ describe("after_tool_call fires exactly once in embedded runs", () => {
       blocked: false,
       params,
     }));
-    await loadFreshAfterToolCallModulesForTest();
   });
 
   function resolveAdapterDefinition(tool: Parameters<typeof toToolDefinitions>[0][number]) {

src/agents/pi-tools.read.workspace-root-guard.test.ts

@@ -1,5 +1,5 @@
 import path from "node:path";
-import { beforeEach, describe, expect, it, vi } from "vitest";
+import { beforeAll, beforeEach, describe, expect, it, vi } from "vitest";
 import type { AnyAgentTool } from "./pi-tools.types.js";
 
 const mocks = vi.hoisted(() => ({
@@ -24,19 +24,21 @@ function createToolHarness() {
 }
 
 async function loadModule() {
-  return await import("./pi-tools.read.js");
+  ({ wrapToolWorkspaceRootGuardWithOptions } = await import("./pi-tools.read.js"));
 }
 
+let wrapToolWorkspaceRootGuardWithOptions: typeof import("./pi-tools.read.js").wrapToolWorkspaceRootGuardWithOptions;
+
 describe("wrapToolWorkspaceRootGuardWithOptions", () => {
   const root = "/tmp/root";
 
+  beforeAll(loadModule);
+
   beforeEach(() => {
     mocks.assertSandboxPath.mockClear();
-    vi.resetModules();
   });
 
   it("maps container workspace paths to host workspace root", async () => {
-    const { wrapToolWorkspaceRootGuardWithOptions } = await loadModule();
     const { tool } = createToolHarness();
     const wrapped = wrapToolWorkspaceRootGuardWithOptions(tool, root, {
       containerWorkdir: "/workspace",
@@ -52,7 +54,6 @@ describe("wrapToolWorkspaceRootGuardWithOptions", () => {
   });
 
   it("maps file:// container workspace paths to host workspace root", async () => {
-    const { wrapToolWorkspaceRootGuardWithOptions } = await loadModule();
     const { tool } = createToolHarness();
     const wrapped = wrapToolWorkspaceRootGuardWithOptions(tool, root, {
       containerWorkdir: "/workspace",
@@ -68,7 +69,6 @@ describe("wrapToolWorkspaceRootGuardWithOptions", () => {
   });
 
   it("does not remap remote-host file:// paths", async () => {
-    const { wrapToolWorkspaceRootGuardWithOptions } = await loadModule();
     const { tool } = createToolHarness();
     const wrapped = wrapToolWorkspaceRootGuardWithOptions(tool, root, {
       containerWorkdir: "/workspace",
@@ -84,7 +84,6 @@ describe("wrapToolWorkspaceRootGuardWithOptions", () => {
   });
 
   it("maps @-prefixed container workspace paths to host workspace root", async () => {
-    const { wrapToolWorkspaceRootGuardWithOptions } = await loadModule();
     const { tool } = createToolHarness();
     const wrapped = wrapToolWorkspaceRootGuardWithOptions(tool, root, {
       containerWorkdir: "/workspace",
@@ -100,7 +99,6 @@ describe("wrapToolWorkspaceRootGuardWithOptions", () => {
   });
 
   it("normalizes @-prefixed absolute paths before guard checks", async () => {
-    const { wrapToolWorkspaceRootGuardWithOptions } = await loadModule();
     const { tool } = createToolHarness();
     const wrapped = wrapToolWorkspaceRootGuardWithOptions(tool, root, {
       containerWorkdir: "/workspace",
@@ -116,7 +114,6 @@ describe("wrapToolWorkspaceRootGuardWithOptions", () => {
   });
 
   it("does not remap absolute paths outside the configured container workdir", async () => {
-    const { wrapToolWorkspaceRootGuardWithOptions } = await loadModule();
     const { tool } = createToolHarness();
     const wrapped = wrapToolWorkspaceRootGuardWithOptions(tool, root, {
       containerWorkdir: "/workspace",

src/agents/subagent-registry.steer-restart.test.ts

@@ -1,4 +1,4 @@
-import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import { afterEach, beforeAll, beforeEach, describe, expect, it, vi } from "vitest";
 
 const noop = () => {};
 let lifecycleHandler:
@@ -101,12 +101,15 @@ describe("subagent registry steer restarts", () => {
   const MAIN_REQUESTER_SESSION_KEY = "agent:main:main";
   const MAIN_REQUESTER_DISPLAY_KEY = "main";
 
-  beforeEach(async () => {
-    vi.resetModules();
-    lifecycleHandler = undefined;
+  beforeAll(async () => {
     mod = await import("./subagent-registry.js");
   });
 
+  beforeEach(() => {
+    lifecycleHandler = undefined;
+    mod.resetSubagentRegistryForTests({ persist: false });
+  });
+
   const flushAnnounce = async () => {
     await new Promise<void>((resolve) => setImmediate(resolve));
   };

src/agents/tools/sessions.test.ts

@@ -57,7 +57,6 @@ type SessionsListResult = Awaited<
 >;
 
 beforeAll(async () => {
-  vi.resetModules();
   ({ createSessionsListTool } = await import("./sessions-list-tool.js"));
   ({ createSessionsSendTool } = await import("./sessions-send-tool.js"));
   ({ resolveAnnounceTarget } = await import("./sessions-announce-target.js"));