Fix tests to wrap all calls changing the UI with act. (#12268)

Author: jacob314

packages/a2a-server/vitest.config.ts

@@ -11,7 +11,6 @@ export default defineConfig({
   test: {
     include: ['**/*.{test,spec}.?(c|m)[jt]s?(x)'],
     exclude: ['**/node_modules/**', '**/dist/**'],
-    environment: 'jsdom',
     globals: true,
     reporters: ['default', 'junit'],
     silent: true,

packages/cli/src/test-utils/async.ts

@@ -0,0 +1,34 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { act } from 'react';
+
+// The waitFor from vitest doesn't properly wrap in act(), so we have to
+// implement our own like the one in @testing-library/react
+// or @testing-library/react-native
+// The version of waitFor from vitest is still fine to use if you aren't waiting
+// for React state updates.
+export async function waitFor(
+  assertion: () => void,
+  { timeout = 1000, interval = 50 } = {},
+): Promise<void> {
+  const startTime = Date.now();
+
+  while (true) {
+    try {
+      assertion();
+      return;
+    } catch (error) {
+      if (Date.now() - startTime > timeout) {
+        throw error;
+      }
+
+      await act(async () => {
+        await new Promise((resolve) => setTimeout(resolve, interval));
+      });
+    }
+  }
+}

packages/cli/src/test-utils/render.test.tsx

@@ -6,7 +6,37 @@
 
 import { describe, it, expect, vi } from 'vitest';
 import { useState, useEffect } from 'react';
-import { renderHook } from './render.js';
+import { Text } from 'ink';
+import { renderHook, render } from './render.js';
+import { waitFor } from './async.js';
+
+describe('render', () => {
+  it('should render a component', () => {
+    const { lastFrame } = render(<Text>Hello World</Text>);
+    expect(lastFrame()).toBe('Hello World');
+  });
+
+  it('should support rerender', () => {
+    const { lastFrame, rerender } = render(<Text>Hello</Text>);
+    expect(lastFrame()).toBe('Hello');
+
+    rerender(<Text>World</Text>);
+    expect(lastFrame()).toBe('World');
+  });
+
+  it('should support unmount', () => {
+    const cleanup = vi.fn();
+    function TestComponent() {
+      useEffect(() => cleanup, []);
+      return <Text>Hello</Text>;
+    }
+
+    const { unmount } = render(<TestComponent />);
+    unmount();
+
+    expect(cleanup).toHaveBeenCalled();
+  });
+});
 
 describe('renderHook', () => {
   it('should rerender with previous props when called without arguments', async () => {
@@ -23,19 +53,19 @@ describe('renderHook', () => {
     });
 
     expect(result.current.value).toBe(1);
-    await vi.waitFor(() => expect(result.current.count).toBe(1));
+    await waitFor(() => expect(result.current.count).toBe(1));
 
     // Rerender with new props
     rerender({ value: 2 });
     expect(result.current.value).toBe(2);
-    await vi.waitFor(() => expect(result.current.count).toBe(2));
+    await waitFor(() => expect(result.current.count).toBe(2));
 
     // Rerender without arguments should use previous props (value: 2)
     // This would previously crash or pass undefined if not fixed
     rerender();
     expect(result.current.value).toBe(2);
     // Count should not increase because value didn't change
-    await vi.waitFor(() => expect(result.current.count).toBe(2));
+    await waitFor(() => expect(result.current.count).toBe(2));
   });
 
   it('should handle initial render without props', () => {

packages/cli/src/test-utils/render.tsx

@@ -4,7 +4,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { render } from 'ink-testing-library';
+import { render as inkRender } from 'ink-testing-library';
 import type React from 'react';
 import { act } from 'react';
 import { LoadedSettings, type Settings } from '../config/settings.js';
@@ -19,6 +19,34 @@ import { VimModeProvider } from '../ui/contexts/VimModeContext.js';
 
 import { type Config } from '@google/gemini-cli-core';
 
+// Wrapper around ink-testing-library's render that ensures act() is called
+export const render = (
+  tree: React.ReactElement,
+): ReturnType<typeof inkRender> => {
+  let renderResult: ReturnType<typeof inkRender> =
+    undefined as unknown as ReturnType<typeof inkRender>;
+  act(() => {
+    renderResult = inkRender(tree);
+  });
+
+  const originalUnmount = renderResult.unmount;
+  const originalRerender = renderResult.rerender;
+
+  return {
+    ...renderResult,
+    unmount: () => {
+      act(() => {
+        originalUnmount();
+      });
+    },
+    rerender: (newTree: React.ReactElement) => {
+      act(() => {
+        originalRerender(newTree);
+      });
+    },
+  };
+};
+
 const mockConfig = {
   getModel: () => 'gemini-pro',
   getTargetDir: () =>

packages/cli/src/ui/App.test.tsx

@@ -5,7 +5,7 @@
  */
 
 import { describe, it, expect, vi, type Mock } from 'vitest';
-import { render } from 'ink-testing-library';
+import { render } from '../test-utils/render.js';
 import { Text, useIsScreenReaderEnabled } from 'ink';
 import { makeFakeConfig } from '@google/gemini-cli-core';
 import { App } from './App.js';