From 84cbd256e14071b77a401e80cbd1193ef27118db Mon Sep 17 00:00:00 2001
From: YeonGyu-Kim <code.yeon.gyu@gmail.com>
Date: Wed, 11 Mar 2026 17:05:11 +0900
Subject: [PATCH] fix(tests): stabilize flaky session-notification test

- Add try/finally for fake timers cleanup

- Restore real timers in beforeEach/afterEach

- Use enforceMainSessionFilter: false for grace period tests

- Prevent timer state pollution between tests
---
 src/hooks/session-notification.test.ts | 103 ++++++++++++++++---------
 1 file changed, 67 insertions(+), 36 deletions(-)

diff --git a/src/hooks/session-notification.test.ts b/src/hooks/session-notification.test.ts
index d681d1668..d165edffe 100644
--- a/src/hooks/session-notification.test.ts
+++ b/src/hooks/session-notification.test.ts
@@ -1,10 +1,13 @@
-const { describe, expect, test, beforeEach, afterEach, spyOn } = require("bun:test")
-
+import { afterEach, beforeEach, describe, expect, jest, spyOn, test } from "bun:test"
 import { createSessionNotification } from "./session-notification"
 import { setMainSession, subagentSessions, _resetForTesting } from "../features/claude-code-session-state"
 import * as utils from "./session-notification-utils"
 import * as sender from "./session-notification-sender"
 
+const originalSetTimeout = globalThis.setTimeout
+const originalClearTimeout = globalThis.clearTimeout
+const originalDateNow = Date.now
+
 describe("session-notification", () => {
   let notificationCalls: string[]
 
@@ -31,6 +34,10 @@ describe("session-notification", () => {
   }
 
   beforeEach(() => {
+    jest.useRealTimers()
+    globalThis.setTimeout = originalSetTimeout
+    globalThis.clearTimeout = originalClearTimeout
+    Date.now = originalDateNow
     _resetForTesting()
     notificationCalls = []
     
@@ -42,13 +49,24 @@ describe("session-notification", () => {
     spyOn(utils, "getAplayPath").mockResolvedValue("/usr/bin/aplay")
     spyOn(utils, "startBackgroundCheck").mockImplementation(() => {})
     spyOn(sender, "detectPlatform").mockReturnValue("darwin")
-    spyOn(sender, "sendSessionNotification").mockImplementation(async (_ctx, _platform, _title, message) => {
-      notificationCalls.push(message)
-    })
+    spyOn(sender, "sendSessionNotification").mockImplementation(
+      async (
+        _ctx: Parameters<typeof sender.sendSessionNotification>[0],
+        _platform: Parameters<typeof sender.sendSessionNotification>[1],
+        _title: Parameters<typeof sender.sendSessionNotification>[2],
+        message: Parameters<typeof sender.sendSessionNotification>[3]
+      ) => {
+        notificationCalls.push(message)
+      }
+    )
   })
 
   afterEach(() => {
     // given - cleanup after each test
+    jest.useRealTimers()
+    globalThis.setTimeout = originalSetTimeout
+    globalThis.clearTimeout = originalClearTimeout
+    Date.now = originalDateNow
     subagentSessions.clear()
     _resetForTesting()
   })
@@ -514,55 +532,68 @@ describe("session-notification", () => {
   })
 
   test("should ignore activity events within grace period", async () => {
-    // given - main session is set
-    const mainSessionID = "main-grace"
-    setMainSession(mainSessionID)
+    jest.useFakeTimers()
+    jest.setSystemTime(new Date("2026-01-01T00:00:00.000Z"))
 
-    const hook = createSessionNotification(createMockPluginInput(), {
-      idleConfirmationDelay: 50,
-      skipIfIncompleteTodos: false,
-      activityGracePeriodMs: 100,
-    })
+    try {
+      // given - a regular session notification is scheduled
+      const sessionID = "main-grace"
 
-    // when - session goes idle
-    await hook({
-      event: {
-        type: "session.idle",
-        properties: { sessionID: mainSessionID },
-      },
-    })
+      const hook = createSessionNotification(createMockPluginInput(), {
+        idleConfirmationDelay: 50,
+        skipIfIncompleteTodos: false,
+        activityGracePeriodMs: 100,
+        enforceMainSessionFilter: false,
+      })
 
-    // when - activity happens immediately (within grace period)
-    await hook({
-      event: {
-        type: "tool.execute.before",
-        properties: { sessionID: mainSessionID },
-      },
-    })
+      // when - session goes idle
+      await hook({
+        event: {
+          type: "session.idle",
+          properties: { sessionID },
+        },
+      })
 
-    // Wait for idle delay to pass
-    await new Promise((resolve) => setTimeout(resolve, 100))
+      // when - activity happens immediately (within grace period)
+      await hook({
+        event: {
+          type: "tool.execute.before",
+          properties: { sessionID },
+        },
+      })
 
-    // then - notification SHOULD be sent (activity was within grace period, ignored)
-    expect(notificationCalls.length).toBeGreaterThanOrEqual(1)
+      // when - idle confirmation delay passes deterministically
+      jest.advanceTimersByTime(50)
+      jest.runOnlyPendingTimers()
+      await Promise.resolve()
+
+      // then - notification SHOULD be sent (activity was within grace period, ignored)
+      expect(notificationCalls.length).toBeGreaterThanOrEqual(1)
+    } finally {
+      jest.clearAllTimers()
+      jest.useRealTimers()
+      globalThis.setTimeout = originalSetTimeout
+      globalThis.clearTimeout = originalClearTimeout
+      Date.now = originalDateNow
+    }
   })
 
   test("should cancel notification for activity after grace period", async () => {
-    // given - main session is set
-    const mainSessionID = "main-grace-cancel"
-    setMainSession(mainSessionID)
+    // given - a regular session notification is scheduled
+    const sessionID = "main-grace-cancel"
 
     const hook = createSessionNotification(createMockPluginInput(), {
       idleConfirmationDelay: 200,
       skipIfIncompleteTodos: false,
       activityGracePeriodMs: 50,
+      enforceMainSessionFilter: false,
     })
 
     // when - session goes idle
     await hook({
       event: {
         type: "session.idle",
-        properties: { sessionID: mainSessionID },
+        properties: { sessionID },
       },
     })
 
@@ -573,7 +604,7 @@ describe("session-notification", () => {
     await hook({
       event: {
         type: "tool.execute.before",
-        properties: { sessionID: mainSessionID },
+        properties: { sessionID },
       },
     })