Fix bug where users are unable to re-enter disconnected terminals. (#8765)

packages/cli/src/ui/hooks/useFocus.test.ts

@@ -9,6 +9,8 @@ import { EventEmitter } from 'node:events';
 import { useFocus } from './useFocus.js';
 import { vi } from 'vitest';
 import { useStdin, useStdout } from 'ink';
+import { KeypressProvider } from '../contexts/KeypressContext.js';
+import React from 'react';
 
 // Mock the ink hooks
 vi.mock('ink', async (importOriginal) => {
@@ -23,12 +25,17 @@ vi.mock('ink', async (importOriginal) => {
 const mockedUseStdin = vi.mocked(useStdin);
 const mockedUseStdout = vi.mocked(useStdout);
 
+const wrapper = ({ children }: { children: React.ReactNode }) =>
+  React.createElement(KeypressProvider, null, children);
+
 describe('useFocus', () => {
   let stdin: EventEmitter;
   let stdout: { write: vi.Func };
 
   beforeEach(() => {
     stdin = new EventEmitter();
+    stdin.resume = vi.fn();
+    stdin.pause = vi.fn();
     stdout = { write: vi.fn() };
     mockedUseStdin.mockReturnValue({ stdin } as ReturnType<typeof useStdin>);
     mockedUseStdout.mockReturnValue({ stdout } as unknown as ReturnType<
@@ -38,17 +45,18 @@ describe('useFocus', () => {
 
   afterEach(() => {
     vi.clearAllMocks();
+    stdin.removeAllListeners();
   });
 
   it('should initialize with focus and enable focus reporting', () => {
-    const { result } = renderHook(() => useFocus());
+    const { result } = renderHook(() => useFocus(), { wrapper });
 
     expect(result.current).toBe(true);
     expect(stdout.write).toHaveBeenCalledWith('\x1b[?1004h');
   });
 
   it('should set isFocused to false when a focus-out event is received', () => {
-    const { result } = renderHook(() => useFocus());
+    const { result } = renderHook(() => useFocus(), { wrapper });
 
     // Initial state is focused
     expect(result.current).toBe(true);
@@ -63,7 +71,7 @@ describe('useFocus', () => {
   });
 
   it('should set isFocused to true when a focus-in event is received', () => {
-    const { result } = renderHook(() => useFocus());
+    const { result } = renderHook(() => useFocus(), { wrapper });
 
     // Simulate focus-out to set initial state to false
     act(() => {
@@ -81,20 +89,22 @@ describe('useFocus', () => {
   });
 
   it('should clean up and disable focus reporting on unmount', () => {
-    const { unmount } = renderHook(() => useFocus());
+    const { unmount } = renderHook(() => useFocus(), { wrapper });
 
-    // Ensure listener was attached
-    expect(stdin.listenerCount('data')).toBe(1);
+    // At this point we should have listeners from both KeypressProvider and useFocus
+    const listenerCountAfterMount = stdin.listenerCount('data');
+    expect(listenerCountAfterMount).toBeGreaterThanOrEqual(1);
 
     unmount();
 
     // Assert that the cleanup function was called
     expect(stdout.write).toHaveBeenCalledWith('\x1b[?1004l');
-    expect(stdin.listenerCount('data')).toBe(0);
+    // Ensure useFocus listener was removed (but KeypressProvider listeners may remain)
+    expect(stdin.listenerCount('data')).toBeLessThan(listenerCountAfterMount);
   });
 
   it('should handle multiple focus events correctly', () => {
-    const { result } = renderHook(() => useFocus());
+    const { result } = renderHook(() => useFocus(), { wrapper });
 
     act(() => {
       stdin.emit('data', Buffer.from('\x1b[O'));
@@ -116,4 +126,20 @@ describe('useFocus', () => {
     });
     expect(result.current).toBe(true);
   });
+
+  it('restores focus on keypress after focus is lost', () => {
+    const { result } = renderHook(() => useFocus(), { wrapper });
+
+    // Simulate focus-out event
+    act(() => {
+      stdin.emit('data', Buffer.from('\x1b[O'));
+    });
+    expect(result.current).toBe(false);
+
+    // Simulate a keypress
+    act(() => {
+      stdin.emit('data', Buffer.from('a'));
+    });
+    expect(result.current).toBe(true);
+  });
 });