fix shutdown manager tests

Author: nicktrn

packages/core/src/v3/serverOnly/shutdownManager.test.ts

@@ -1,57 +1,55 @@
 import { describe, test, expect, vi, beforeEach } from "vitest";
-import { shutdownManager } from "./shutdownManager.js";
-
-// Type assertion to access private members for testing
-type PrivateShutdownManager = {
-  handlers: Map<string, { handler: NodeJS.SignalsListener; signals: Array<"SIGTERM" | "SIGINT"> }>;
-  shutdown: (signal: "SIGTERM" | "SIGINT") => Promise<void>;
-  _reset: () => void;
-};
+import { ShutdownManager } from "./shutdownManager.js";
 
 describe("ShutdownManager", { concurrent: false }, () => {
-  const manager = shutdownManager as unknown as PrivateShutdownManager;
   // Mock process.exit to prevent actual exit
   const mockExit = vi.spyOn(process, "exit").mockImplementation(() => undefined as never);
 
   beforeEach(() => {
-    // Clear all mocks and reset the manager before each test
     vi.clearAllMocks();
-    manager._reset();
   });
 
   test("should successfully register a new handler", () => {
+    const manager = new ShutdownManager(false);
+
     const handler = vi.fn();
-    shutdownManager.register("test-handler", handler);
+    manager.register("test-handler", handler);
 
-    expect(manager.handlers.has("test-handler")).toBe(true);
-    const registeredHandler = manager.handlers.get("test-handler");
+    expect(manager._getHandlersForTesting().has("test-handler")).toBe(true);
+    const registeredHandler = manager._getHandlersForTesting().get("test-handler");
     expect(registeredHandler?.handler).toBe(handler);
     expect(registeredHandler?.signals).toEqual(["SIGTERM", "SIGINT"]);
   });
 
   test("should throw error when registering duplicate handler name", () => {
+    const manager = new ShutdownManager(false);
+
     const handler = vi.fn();
-    shutdownManager.register("duplicate-handler", handler);
+    manager.register("duplicate-handler", handler);
 
     expect(() => {
-      shutdownManager.register("duplicate-handler", handler);
+      manager.register("duplicate-handler", handler);
     }).toThrow('Shutdown handler "duplicate-handler" already registered');
   });
 
   test("should register handler with custom signals", () => {
+    const manager = new ShutdownManager(false);
+
     const handler = vi.fn();
-    shutdownManager.register("custom-signals", handler, ["SIGTERM"]);
+    manager.register("custom-signals", handler, ["SIGTERM"]);
 
-    const registeredHandler = manager.handlers.get("custom-signals");
+    const registeredHandler = manager._getHandlersForTesting().get("custom-signals");
     expect(registeredHandler?.signals).toEqual(["SIGTERM"]);
   });
 
   test("should call registered handlers when shutdown is triggered", async () => {
+    const manager = new ShutdownManager(false);
+
     const handler1 = vi.fn();
     const handler2 = vi.fn();
 
-    shutdownManager.register("handler1", handler1);
-    shutdownManager.register("handler2", handler2);
+    manager.register("handler1", handler1);
+    manager.register("handler2", handler2);
 
     await manager.shutdown("SIGTERM");
 
@@ -61,11 +59,13 @@ describe("ShutdownManager", { concurrent: false }, () => {
   });
 
   test("should only call handlers registered for specific signal", async () => {
+    const manager = new ShutdownManager(false);
+
     const handler1 = vi.fn();
     const handler2 = vi.fn();
 
-    shutdownManager.register("handler1", handler1, ["SIGTERM"]);
-    shutdownManager.register("handler2", handler2, ["SIGINT"]);
+    manager.register("handler1", handler1, ["SIGTERM"]);
+    manager.register("handler2", handler2, ["SIGINT"]);
 
     await manager.shutdown("SIGTERM");
 
@@ -75,11 +75,13 @@ describe("ShutdownManager", { concurrent: false }, () => {
   });
 
   test("should handle errors in shutdown handlers gracefully", async () => {
+    const manager = new ShutdownManager(false);
+
     const successHandler = vi.fn();
     const errorHandler = vi.fn().mockRejectedValue(new Error("Handler failed"));
 
-    shutdownManager.register("success-handler", successHandler);
-    shutdownManager.register("error-handler", errorHandler);
+    manager.register("success-handler", successHandler);
+    manager.register("error-handler", errorHandler);
 
     await manager.shutdown("SIGTERM");
 
@@ -89,8 +91,10 @@ describe("ShutdownManager", { concurrent: false }, () => {
   });
 
   test("should only run shutdown sequence once even if called multiple times", async () => {
+    const manager = new ShutdownManager(false);
+
     const handler = vi.fn();
-    shutdownManager.register("test-handler", handler);
+    manager.register("test-handler", handler);
 
     await Promise.all([manager.shutdown("SIGTERM"), manager.shutdown("SIGTERM")]);
 
@@ -99,23 +103,27 @@ describe("ShutdownManager", { concurrent: false }, () => {
     expect(mockExit).toHaveBeenCalledWith(128 + 15);
   });
 
-  test("should exit with correct signal number", async () => {
-    const handler = vi.fn();
-    shutdownManager.register("test-handler", handler);
+  test("should exit with correct signal number on SIGINT", async () => {
+    const manager = new ShutdownManager(false);
+
+    manager.register("test-handler", vi.fn());
 
     await manager.shutdown("SIGINT");
     expect(mockExit).toHaveBeenCalledWith(128 + 2); // SIGINT number
+  });
 
-    vi.clearAllMocks();
-    manager._reset();
-    shutdownManager.register("test-handler", handler);
+  test("should exit with correct signal number on SIGTERM", async () => {
+    const manager = new ShutdownManager(false);
+
+    manager.register("test-handler", vi.fn());
 
     await manager.shutdown("SIGTERM");
     expect(mockExit).toHaveBeenCalledWith(128 + 15); // SIGTERM number
   });
 
   test("should only exit after all handlers have finished", async () => {
     const sequence: string[] = [];
+    const manager = new ShutdownManager(false);
 
     const handler1 = vi.fn().mockImplementation(async () => {
       sequence.push("handler1 start");
@@ -144,9 +152,9 @@ describe("ShutdownManager", { concurrent: false }, () => {
       return undefined as never;
     });
 
-    shutdownManager.register("handler1", handler1);
-    shutdownManager.register("handler2", handler2);
-    shutdownManager.register("handler3", handler3);
+    manager.register("handler1", handler1);
+    manager.register("handler2", handler2);
+    manager.register("handler3", handler3);
 
     await manager.shutdown("SIGTERM");
 

packages/core/src/v3/serverOnly/shutdownManager.ts

@@ -1,11 +1,12 @@
+import { isTest } from "std-env";
 import { SimpleStructuredLogger } from "../utils/structuredLogger.js";
 import { singleton } from "./singleton.js";
 
 type ShutdownHandler = NodeJS.SignalsListener;
 // We intentionally keep these limited to avoid unexpected issues with signal handling
 type ShutdownSignal = Extract<NodeJS.Signals, "SIGTERM" | "SIGINT">;
 
-class ShutdownManager {
+export class ShutdownManager {
   private isShuttingDown = false;
   private signalNumbers: Record<ShutdownSignal, number> = {
     SIGINT: 2,
@@ -16,7 +17,9 @@ class ShutdownManager {
   private handlers: Map<string, { handler: ShutdownHandler; signals: ShutdownSignal[] }> =
     new Map();
 
-  constructor() {
+  constructor(private disableForTesting = true) {
+    if (disableForTesting) return;
+
     process.on("SIGTERM", () => this.shutdown("SIGTERM"));
     process.on("SIGINT", () => this.shutdown("SIGINT"));
   }
@@ -26,13 +29,17 @@ class ShutdownManager {
     handler: ShutdownHandler,
     signals: ShutdownSignal[] = ["SIGTERM", "SIGINT"]
   ) {
+    if (!this.isEnabled()) return;
+
     if (this.handlers.has(name)) {
       throw new Error(`Shutdown handler "${name}" already registered`);
     }
     this.handlers.set(name, { handler, signals });
   }
 
   unregister(name: string) {
+    if (!this.isEnabled()) return;
+
     if (!this.handlers.has(name)) {
       throw new Error(`Shutdown handler "${name}" not registered`);
     }
@@ -41,6 +48,8 @@ class ShutdownManager {
   }
 
   async shutdown(signal: ShutdownSignal) {
+    if (!this.isEnabled()) return;
+
     if (this.isShuttingDown) return;
     this.isShuttingDown = true;
 
@@ -83,10 +92,20 @@ class ShutdownManager {
     }
   }
 
-  // For testing purposes only - keep this
-  private _reset() {
-    this.isShuttingDown = false;
-    this.handlers.clear();
+  private isEnabled() {
+    if (!this.disableForTesting) {
+      return true;
+    }
+
+    return !isTest;
+  }
+
+  // Only for testing
+  public _getHandlersForTesting(): ReadonlyMap<
+    string,
+    { handler: ShutdownHandler; signals: ShutdownSignal[] }
+  > {
+    return new Map(this.handlers);
   }
 }
 

packages/redis-worker/src/worker.test.ts

@@ -1,17 +1,12 @@
 import { redisTest } from "@internal/testcontainers";
 import { Logger } from "@trigger.dev/core/logger";
 import { describe } from "node:test";
-import { afterEach, expect } from "vitest";
+import { expect } from "vitest";
 import { z } from "zod";
 import { Worker } from "./worker.js";
 import { createRedisClient } from "@internal/redis";
-import { shutdownManager } from "@trigger.dev/core/v3/serverOnly";
 
 describe("Worker", () => {
-  afterEach(() => {
-    (shutdownManager as unknown as { _reset: () => void })._reset();
-  });
-
   redisTest("Process items that don't throw", { timeout: 30_000 }, async ({ redisContainer }) => {
     const processedItems: number[] = [];
     const worker = new Worker({