Migrate tests to use avoid jsdom (#12118)

packages/cli/src/ui/contexts/KeypressContext.test.tsx

@@ -4,10 +4,9 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-/** @vitest-environment jsdom */
-
 import type React from 'react';
-import { renderHook, act, waitFor } from '@testing-library/react';
+import { act } from 'react';
+import { renderHook } from '../../test-utils/render.js';
 import type { Mock } from 'vitest';
 import { vi } from 'vitest';
 import type { Key } from './KeypressContext.js';
@@ -370,7 +369,7 @@ describe('KeypressContext - Kitty Protocol', () => {
         stdin.write(PASTE_END);
       });
 
-      await waitFor(() => {
+      await vi.waitFor(() => {
         // Expect the handler to be called exactly once for the entire paste
         expect(keyHandler).toHaveBeenCalledTimes(1);
       });
@@ -399,7 +398,7 @@ describe('KeypressContext - Kitty Protocol', () => {
         stdin.write(PASTE_END);
       });
 
-      await waitFor(() => {
+      await vi.waitFor(() => {
         expect(keyHandler).toHaveBeenCalledTimes(1);
       });
 
@@ -427,7 +426,7 @@ describe('KeypressContext - Kitty Protocol', () => {
         stdin.write(PASTE_END.slice(3));
       });
 
-      await waitFor(() => {
+      await vi.waitFor(() => {
         expect(keyHandler).toHaveBeenCalledTimes(1);
       });
 
@@ -1193,7 +1192,7 @@ describe('Kitty Sequence Parsing', () => {
     }
 
     // Should parse once complete
-    await waitFor(() => {
+    await vi.waitFor(() => {
       expect(keyHandler).toHaveBeenCalledWith(
         expect.objectContaining({
           name: 'escape',

packages/cli/src/ui/contexts/SessionContext.test.tsx

@@ -4,17 +4,40 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-/** @vitest-environment jsdom */
-
-import { type MutableRefObject } from 'react';
+import { type MutableRefObject, Component, type ReactNode } from 'react';
 import { render } from 'ink-testing-library';
-import { renderHook } from '@testing-library/react';
-import { act } from 'react-dom/test-utils';
+
+import { act } from 'react';
 import type { SessionMetrics } from './SessionContext.js';
 import { SessionStatsProvider, useSessionStats } from './SessionContext.js';
 import { describe, it, expect, vi } from 'vitest';
 import { uiTelemetryService } from '@google/gemini-cli-core';
 
+class ErrorBoundary extends Component<
+  { children: ReactNode; onError: (error: Error) => void },
+  { hasError: boolean }
+> {
+  constructor(props: { children: ReactNode; onError: (error: Error) => void }) {
+    super(props);
+    this.state = { hasError: false };
+  }
+
+  static getDerivedStateFromError(_error: Error) {
+    return { hasError: true };
+  }
+
+  override componentDidCatch(error: Error) {
+    this.props.onError(error);
+  }
+
+  override render() {
+    if (this.state.hasError) {
+      return null;
+    }
+    return this.props.children;
+  }
+}
+
 /**
  * A test harness component that uses the hook and exposes the context value
  * via a mutable ref. This allows us to interact with the context's functions
@@ -208,16 +231,22 @@ describe('SessionStatsContext', () => {
   });
 
   it('should throw an error when useSessionStats is used outside of a provider', () => {
-    // Suppress console.error for this test since we expect an error
+    const onError = vi.fn();
+    // Suppress console.error from React for this test
     const consoleSpy = vi.spyOn(console, 'error').mockImplementation(() => {});
 
-    try {
-      // Expect renderHook itself to throw when the hook is used outside a provider
-      expect(() => {
-        renderHook(() => useSessionStats());
-      }).toThrow('useSessionStats must be used within a SessionStatsProvider');
-    } finally {
-      consoleSpy.mockRestore();
-    }
+    render(
+      <ErrorBoundary onError={onError}>
+        <TestHarness contextRef={{ current: undefined }} />
+      </ErrorBoundary>,
+    );
+
+    expect(onError).toHaveBeenCalledWith(
+      expect.objectContaining({
+        message: 'useSessionStats must be used within a SessionStatsProvider',
+      }),
+    );
+
+    consoleSpy.mockRestore();
   });
 });