Tests: share partial module mock helper

src/agents/acp-spawn-parent-stream.test.ts

@@ -1,4 +1,5 @@
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import { mergeMockedModule } from "../test-utils/vitest-module-mocks.js";
 
 const enqueueSystemEventMock = vi.fn();
 const requestHeartbeatNowMock = vi.fn();
@@ -10,14 +11,13 @@ vi.mock("../infra/system-events.js", () => ({
   enqueueSystemEvent: (...args: unknown[]) => enqueueSystemEventMock(...args),
 }));
 
-vi.mock("../infra/heartbeat-wake.js", async () => {
-  const actual = await vi.importActual<typeof import("../infra/heartbeat-wake.js")>(
-    "../infra/heartbeat-wake.js",
+vi.mock("../infra/heartbeat-wake.js", async (importOriginal) => {
+  return await mergeMockedModule(
+    await importOriginal<typeof import("../infra/heartbeat-wake.js")>(),
+    () => ({
+      requestHeartbeatNow: (...args: unknown[]) => requestHeartbeatNowMock(...args),
+    }),
   );
-  return {
-    ...actual,
-    requestHeartbeatNow: (...args: unknown[]) => requestHeartbeatNowMock(...args),
-  };
 });
 
 vi.mock("../acp/runtime/session-meta.js", () => ({
@@ -39,13 +39,14 @@ async function loadFreshAcpSpawnParentStreamModulesForTest() {
     enqueueSystemEvent: (...args: unknown[]) => enqueueSystemEventMock(...args),
   }));
   vi.doMock("../infra/heartbeat-wake.js", async () => {
-    const actual = await vi.importActual<typeof import("../infra/heartbeat-wake.js")>(
-      "../infra/heartbeat-wake.js",
+    return await mergeMockedModule(
+      await vi.importActual<typeof import("../infra/heartbeat-wake.js")>(
+        "../infra/heartbeat-wake.js",
+      ),
+      () => ({
+        requestHeartbeatNow: (...args: unknown[]) => requestHeartbeatNowMock(...args),
+      }),
     );
-    return {
-      ...actual,
-      requestHeartbeatNow: (...args: unknown[]) => requestHeartbeatNowMock(...args),
-    };
   });
   vi.doMock("../acp/runtime/session-meta.js", () => ({
     readAcpSessionEntry: (...args: unknown[]) => readAcpSessionEntryMock(...args),

src/agents/bash-tools.exec-runtime.test.ts

@@ -1,4 +1,5 @@
 import { beforeEach, describe, expect, it, vi } from "vitest";
+import { mergeMockedModule } from "../test-utils/vitest-module-mocks.js";
 
 const requestHeartbeatNowMock = vi.hoisted(() => vi.fn());
 const enqueueSystemEventMock = vi.hoisted(() => vi.fn());
@@ -46,13 +47,14 @@ describe("emitExecSystemEvent", () => {
     requestHeartbeatNowMock.mockClear();
     enqueueSystemEventMock.mockClear();
     vi.doMock("../infra/heartbeat-wake.js", async () => {
-      const actual = await vi.importActual<typeof import("../infra/heartbeat-wake.js")>(
-        "../infra/heartbeat-wake.js",
+      return await mergeMockedModule(
+        await vi.importActual<typeof import("../infra/heartbeat-wake.js")>(
+          "../infra/heartbeat-wake.js",
+        ),
+        () => ({
+          requestHeartbeatNow: requestHeartbeatNowMock,
+        }),
       );
-      return {
-        ...actual,
-        requestHeartbeatNow: requestHeartbeatNowMock,
-      };
     });
     vi.doMock("../infra/system-events.js", () => ({
       enqueueSystemEvent: enqueueSystemEventMock,

src/agents/cli-runner.test-support.ts

@@ -6,6 +6,7 @@ import { buildOpenAICodexCliBackend } from "../../extensions/openai/test-api.js"
 import type { OpenClawConfig } from "../config/config.js";
 import { createEmptyPluginRegistry } from "../plugins/registry.js";
 import { setActivePluginRegistry } from "../plugins/runtime.js";
+import { mergeMockedModule } from "../test-utils/vitest-module-mocks.js";
 import type { EmbeddedContextFile } from "./pi-embedded-helpers.js";
 import type { WorkspaceBootstrapFile } from "./workspace.js";
 
@@ -43,14 +44,13 @@ vi.mock("../infra/system-events.js", () => ({
   enqueueSystemEvent: (...args: unknown[]) => enqueueSystemEventMock(...args),
 }));
 
-vi.mock("../infra/heartbeat-wake.js", async () => {
-  const actual = await vi.importActual<typeof import("../infra/heartbeat-wake.js")>(
-    "../infra/heartbeat-wake.js",
+vi.mock("../infra/heartbeat-wake.js", async (importOriginal) => {
+  return await mergeMockedModule(
+    await importOriginal<typeof import("../infra/heartbeat-wake.js")>(),
+    () => ({
+      requestHeartbeatNow: (...args: unknown[]) => requestHeartbeatNowMock(...args),
+    }),
   );
-  return {
-    ...actual,
-    requestHeartbeatNow: (...args: unknown[]) => requestHeartbeatNowMock(...args),
-  };
 });
 
 vi.mock("./bootstrap-files.js", () => ({
@@ -166,13 +166,14 @@ export async function setupCliRunnerTestModule() {
     enqueueSystemEvent: (...args: unknown[]) => enqueueSystemEventMock(...args),
   }));
   vi.doMock("../infra/heartbeat-wake.js", async () => {
-    const actual = await vi.importActual<typeof import("../infra/heartbeat-wake.js")>(
-      "../infra/heartbeat-wake.js",
+    return await mergeMockedModule(
+      await vi.importActual<typeof import("../infra/heartbeat-wake.js")>(
+        "../infra/heartbeat-wake.js",
+      ),
+      () => ({
+        requestHeartbeatNow: (...args: unknown[]) => requestHeartbeatNowMock(...args),
+      }),
     );
-    return {
-      ...actual,
-      requestHeartbeatNow: (...args: unknown[]) => requestHeartbeatNowMock(...args),
-    };
   });
   vi.doMock("./bootstrap-files.js", () => ({
     makeBootstrapWarn: () => () => {},

src/gateway/server-cron.test.ts

@@ -4,6 +4,7 @@ import { beforeEach, describe, expect, it, vi } from "vitest";
 import type { CliDeps } from "../cli/deps.js";
 import type { OpenClawConfig } from "../config/config.js";
 import { SsrFBlockedError } from "../infra/net/ssrf.js";
+import { mergeMockedModule } from "../test-utils/vitest-module-mocks.js";
 
 const {
   enqueueSystemEventMock,
@@ -31,14 +32,13 @@ vi.mock("../infra/system-events.js", () => ({
   enqueueSystemEvent,
 }));
 
-vi.mock("../infra/heartbeat-wake.js", async () => {
-  const actual = await vi.importActual<typeof import("../infra/heartbeat-wake.js")>(
-    "../infra/heartbeat-wake.js",
+vi.mock("../infra/heartbeat-wake.js", async (importOriginal) => {
+  return await mergeMockedModule(
+    await importOriginal<typeof import("../infra/heartbeat-wake.js")>(),
+    () => ({
+      requestHeartbeatNow,
+    }),
   );
-  return {
-    ...actual,
-    requestHeartbeatNow,
-  };
 });
 
 vi.mock("../config/config.js", async () => {

src/gateway/server-node-events.test.ts

@@ -1,5 +1,6 @@
 import { beforeEach, describe, expect, it, vi } from "vitest";
 import type { OpenClawConfig } from "../config/config.js";
+import { mergeMockedModule } from "../test-utils/vitest-module-mocks.js";
 import type { loadSessionEntry as loadSessionEntryType } from "./session-utils.js";
 
 const buildSessionLookup = (
@@ -43,18 +44,23 @@ const loadOrCreateDeviceIdentityMock = vi.hoisted(() =>
     privateKeyPem: "private",
   })),
 );
+const normalizeChannelIdMock = vi.hoisted(() =>
+  vi.fn((channel?: string | null) => channel ?? null),
+);
 
 vi.mock("../infra/system-events.js", () => ({
   enqueueSystemEvent: vi.fn(),
 }));
-vi.mock("../infra/heartbeat-wake.js", async () => {
-  const actual = await vi.importActual<typeof import("../infra/heartbeat-wake.js")>(
-    "../infra/heartbeat-wake.js",
+vi.mock("../channels/plugins/index.js", () => ({
+  normalizeChannelId: normalizeChannelIdMock,
+}));
+vi.mock("../infra/heartbeat-wake.js", async (importOriginal) => {
+  return await mergeMockedModule(
+    await importOriginal<typeof import("../infra/heartbeat-wake.js")>(),
+    () => ({
+      requestHeartbeatNow: vi.fn(),
+    }),
   );
-  return {
-    ...actual,
-    requestHeartbeatNow: vi.fn(),
-  };
 });
 vi.mock("../commands/agent.js", () => ({
   agentCommand: ingressAgentCommandMock,
@@ -89,6 +95,7 @@ vi.mock("./session-utils.js", () => ({
   })),
 }));
 
+import { normalizeChannelId } from "../channels/plugins/index.js";
 import type { CliDeps } from "../cli/deps.js";
 import { agentCommand } from "../commands/agent.js";
 import type { HealthSummary } from "../commands/health.js";
@@ -108,6 +115,7 @@ const agentCommandMock = vi.mocked(agentCommand);
 const updateSessionStoreMock = vi.mocked(updateSessionStore);
 const loadSessionEntryMock = vi.mocked(loadSessionEntry);
 const registerApnsRegistrationVi = vi.mocked(registerApnsRegistration);
+const normalizeChannelIdVi = vi.mocked(normalizeChannelId);
 
 function buildCtx(): NodeEventContext {
   return {
@@ -138,6 +146,8 @@ describe("node exec events", () => {
     requestHeartbeatNowMock.mockClear();
     registerApnsRegistrationVi.mockClear();
     loadOrCreateDeviceIdentityMock.mockClear();
+    normalizeChannelIdVi.mockClear();
+    normalizeChannelIdVi.mockImplementation((channel?: string | null) => channel ?? null);
   });
 
   it("enqueues exec.started events", async () => {
@@ -568,6 +578,8 @@ describe("notifications changed events", () => {
     enqueueSystemEventMock.mockClear();
     requestHeartbeatNowMock.mockClear();
     loadSessionEntryMock.mockClear();
+    normalizeChannelIdVi.mockClear();
+    normalizeChannelIdVi.mockImplementation((channel?: string | null) => channel ?? null);
     loadSessionEntryMock.mockImplementation((sessionKey: string) => buildSessionLookup(sessionKey));
     enqueueSystemEventMock.mockReturnValue(true);
   });
@@ -719,6 +731,8 @@ describe("agent request events", () => {
     agentCommandMock.mockClear();
     updateSessionStoreMock.mockClear();
     loadSessionEntryMock.mockClear();
+    normalizeChannelIdVi.mockClear();
+    normalizeChannelIdVi.mockImplementation((channel?: string | null) => channel ?? null);
     agentCommandMock.mockResolvedValue({ status: "ok" } as never);
     updateSessionStoreMock.mockImplementation(async (_storePath, update) => {
       update({});

src/gateway/server-restart-sentinel.test.ts

@@ -1,4 +1,5 @@
 import { beforeEach, describe, expect, it, vi } from "vitest";
+import { mergeMockedModule } from "../test-utils/vitest-module-mocks.js";
 
 const mocks = vi.hoisted(() => ({
   resolveSessionAgentId: vi.fn(() => "agent-from-key"),
@@ -87,14 +88,13 @@ vi.mock("../infra/system-events.js", () => ({
   enqueueSystemEvent: mocks.enqueueSystemEvent,
 }));
 
-vi.mock("../infra/heartbeat-wake.js", async () => {
-  const actual = await vi.importActual<typeof import("../infra/heartbeat-wake.js")>(
-    "../infra/heartbeat-wake.js",
+vi.mock("../infra/heartbeat-wake.js", async (importOriginal) => {
+  return await mergeMockedModule(
+    await importOriginal<typeof import("../infra/heartbeat-wake.js")>(),
+    () => ({
+      requestHeartbeatNow: mocks.requestHeartbeatNow,
+    }),
   );
-  return {
-    ...actual,
-    requestHeartbeatNow: mocks.requestHeartbeatNow,
-  };
 });
 
 vi.mock("../logging/subsystem.js", () => ({

src/test-utils/vitest-module-mocks.ts

@@ -0,0 +1,9 @@
+export async function mergeMockedModule<TModule extends object>(
+  actual: TModule,
+  buildOverrides: (actual: TModule) => Partial<TModule> | Promise<Partial<TModule>>,
+) {
+  return {
+    ...actual,
+    ...(await buildOverrides(actual)),
+  };
+}