1 files changed, 72 insertions, 137 deletions

diff --git a/packages/cli/src/ui/hooks/useConsoleMessages.test.ts b/packages/cli/src/ui/hooks/useConsoleMessages.test.ts
index 3b225ecf..b1d1acd6 100644
--- a/packages/cli/src/ui/hooks/useConsoleMessages.test.ts
+++ b/packages/cli/src/ui/hooks/useConsoleMessages.test.ts
@@ -5,127 +5,105 @@
  */
 
 import { act, renderHook } from '@testing-library/react';
-import { useConsoleMessages } from './useConsoleMessages.js';
-import { ConsoleMessageItem } from '../types.js';
-
-// Mock setTimeout and clearTimeout
-vi.useFakeTimers();
+import { vi } from 'vitest';
+import { useConsoleMessages } from './useConsoleMessages';
+import { useCallback } from 'react';
 
 describe('useConsoleMessages', () => {
-  it('should initialize with an empty array of console messages', () => {
-    const { result } = renderHook(() => useConsoleMessages());
-    expect(result.current.consoleMessages).toEqual([]);
+  beforeEach(() => {
+    vi.useFakeTimers();
   });
 
-  it('should add a new message', () => {
-    const { result } = renderHook(() => useConsoleMessages());
-    const message: ConsoleMessageItem = {
-      type: 'log',
-      content: 'Test message',
-      count: 1,
-    };
-
-    act(() => {
-      result.current.handleNewMessage(message);
-    });
+  afterEach(() => {
+    vi.runOnlyPendingTimers();
+    vi.useRealTimers();
+  });
 
-    act(() => {
-      vi.runAllTimers(); // Process the queue
-    });
+  const useTestableConsoleMessages = () => {
+    const { handleNewMessage, ...rest } = useConsoleMessages();
+    const log = useCallback(
+      (content: string) => handleNewMessage({ type: 'log', content, count: 1 }),
+      [handleNewMessage],
+    );
+    const error = useCallback(
+      (content: string) =>
+        handleNewMessage({ type: 'error', content, count: 1 }),
+      [handleNewMessage],
+    );
+    return {
+      ...rest,
+      log,
+      error,
+      clearConsoleMessages: rest.clearConsoleMessages,
+    };
+  };
 
-    expect(result.current.consoleMessages).toEqual([{ ...message, count: 1 }]);
+  it('should initialize with an empty array of console messages', () => {
+    const { result } = renderHook(() => useTestableConsoleMessages());
+    expect(result.current.consoleMessages).toEqual([]);
   });
 
-  it('should consolidate identical consecutive messages', () => {
-    const { result } = renderHook(() => useConsoleMessages());
-    const message: ConsoleMessageItem = {
-      type: 'log',
-      content: 'Test message',
-      count: 1,
-    };
+  it('should add a new message when log is called', async () => {
+    const { result } = renderHook(() => useTestableConsoleMessages());
 
     act(() => {
-      result.current.handleNewMessage(message);
-      result.current.handleNewMessage(message);
+      result.current.log('Test message');
     });
 
-    act(() => {
-      vi.runAllTimers();
+    await act(async () => {
+      await vi.advanceTimersByTimeAsync(20);
     });
 
-    expect(result.current.consoleMessages).toEqual([{ ...message, count: 2 }]);
+    expect(result.current.consoleMessages).toEqual([
+      { type: 'log', content: 'Test message', count: 1 },
+    ]);
   });
 
-  it('should not consolidate different messages', () => {
-    const { result } = renderHook(() => useConsoleMessages());
-    const message1: ConsoleMessageItem = {
-      type: 'log',
-      content: 'Test message 1',
-      count: 1,
-    };
-    const message2: ConsoleMessageItem = {
-      type: 'error',
-      content: 'Test message 2',
-      count: 1,
-    };
+  it('should batch and count identical consecutive messages', async () => {
+    const { result } = renderHook(() => useTestableConsoleMessages());
 
     act(() => {
-      result.current.handleNewMessage(message1);
-      result.current.handleNewMessage(message2);
+      result.current.log('Test message');
+      result.current.log('Test message');
+      result.current.log('Test message');
     });
 
-    act(() => {
-      vi.runAllTimers();
+    await act(async () => {
+      await vi.advanceTimersByTimeAsync(20);
     });
 
     expect(result.current.consoleMessages).toEqual([
-      { ...message1, count: 1 },
-      { ...message2, count: 1 },
+      { type: 'log', content: 'Test message', count: 3 },
     ]);
   });
 
-  it('should not consolidate messages if type is different', () => {
-    const { result } = renderHook(() => useConsoleMessages());
-    const message1: ConsoleMessageItem = {
-      type: 'log',
-      content: 'Test message',
-      count: 1,
-    };
-    const message2: ConsoleMessageItem = {
-      type: 'error',
-      content: 'Test message',
-      count: 1,
-    };
+  it('should not batch different messages', async () => {
+    const { result } = renderHook(() => useTestableConsoleMessages());
 
     act(() => {
-      result.current.handleNewMessage(message1);
-      result.current.handleNewMessage(message2);
+      result.current.log('First message');
+      result.current.error('Second message');
     });
 
-    act(() => {
-      vi.runAllTimers();
+    await act(async () => {
+      await vi.advanceTimersByTimeAsync(20);
     });
 
     expect(result.current.consoleMessages).toEqual([
-      { ...message1, count: 1 },
-      { ...message2, count: 1 },
+      { type: 'log', content: 'First message', count: 1 },
+      { type: 'error', content: 'Second message', count: 1 },
     ]);
   });
 
-  it('should clear console messages', () => {
-    const { result } = renderHook(() => useConsoleMessages());
-    const message: ConsoleMessageItem = {
-      type: 'log',
-      content: 'Test message',
-      count: 1,
-    };
+  it('should clear all messages when clearConsoleMessages is called', async () => {
+    const { result } = renderHook(() => useTestableConsoleMessages());
 
     act(() => {
-      result.current.handleNewMessage(message);
+      result.current.log('A message');
     });
 
-    act(() => {
-      vi.runAllTimers();
+    await act(async () => {
+      await vi.advanceTimersByTimeAsync(20);
     });
 
     expect(result.current.consoleMessages).toHaveLength(1);
@@ -134,79 +112,36 @@ describe('useConsoleMessages', () => {
       result.current.clearConsoleMessages();
     });
 
-    expect(result.current.consoleMessages).toEqual([]);
+    expect(result.current.consoleMessages).toHaveLength(0);
   });
 
-  it('should clear pending timeout on clearConsoleMessages', () => {
-    const { result } = renderHook(() => useConsoleMessages());
-    const message: ConsoleMessageItem = {
-      type: 'log',
-      content: 'Test message',
-      count: 1,
-    };
+  it('should clear the pending timeout when clearConsoleMessages is called', () => {
+    const { result } = renderHook(() => useTestableConsoleMessages());
+    const clearTimeoutSpy = vi.spyOn(global, 'clearTimeout');
 
     act(() => {
-      result.current.handleNewMessage(message); // This schedules a timeout
+      result.current.log('A message');
     });
 
     act(() => {
       result.current.clearConsoleMessages();
     });
 
-    // Ensure the queue is empty and no more messages are processed
-    act(() => {
-      vi.runAllTimers(); // If timeout wasn't cleared, this would process the queue
-    });
-
-    expect(result.current.consoleMessages).toEqual([]);
+    expect(clearTimeoutSpy).toHaveBeenCalled();
+    clearTimeoutSpy.mockRestore();
   });
 
-  it('should clear message queue on clearConsoleMessages', () => {
-    const { result } = renderHook(() => useConsoleMessages());
-    const message: ConsoleMessageItem = {
-      type: 'log',
-      content: 'Test message',
-      count: 1,
-    };
-
-    act(() => {
-      // Add a message but don't process the queue yet
-      result.current.handleNewMessage(message);
-    });
-
-    act(() => {
-      result.current.clearConsoleMessages();
-    });
-
-    // Process any pending timeouts (should be none related to message queue)
-    act(() => {
-      vi.runAllTimers();
-    });
-
-    // The consoleMessages should be empty because the queue was cleared before processing
-    expect(result.current.consoleMessages).toEqual([]);
-  });
-
-  it('should cleanup timeout on unmount', () => {
-    const { result, unmount } = renderHook(() => useConsoleMessages());
-    const message: ConsoleMessageItem = {
-      type: 'log',
-      content: 'Test message',
-      count: 1,
-    };
+  it('should clean up the timeout on unmount', () => {
+    const { result, unmount } = renderHook(() => useTestableConsoleMessages());
+    const clearTimeoutSpy = vi.spyOn(global, 'clearTimeout');
 
     act(() => {
-      result.current.handleNewMessage(message);
+      result.current.log('A message');
     });
 
     unmount();
 
-    // This is a bit indirect. We check that clearTimeout was called.
-    // If clearTimeout was not called, and we run timers, an error might occur
-    // or the state might change, which it shouldn't after unmount.
-    // Vitest's vi.clearAllTimers() or specific checks for clearTimeout calls
-    // would be more direct if available and easy to set up here.
-    // For now, we rely on the useEffect cleanup pattern.
-    expect(vi.getTimerCount()).toBe(0); // Check if all timers are cleared
+    expect(clearTimeoutSpy).toHaveBeenCalled();
+    clearTimeoutSpy.mockRestore();
   });
 });