bug(core): Fix flaky test by using waitFor. (#5540)

Co-authored-by: Sandy Tao <[email protected]>

1 files changed, 17 insertions, 10 deletions

diff --git a/packages/cli/src/ui/components/InputPrompt.test.tsx b/packages/cli/src/ui/components/InputPrompt.test.tsx
index 2291b5a1..f050ba07 100644
--- a/packages/cli/src/ui/components/InputPrompt.test.tsx
+++ b/packages/cli/src/ui/components/InputPrompt.test.tsx
@@ -5,6 +5,7 @@
  */
 
 import { render } from 'ink-testing-library';
+import { waitFor } from '@testing-library/react';
 import { InputPrompt, InputPromptProps } from './InputPrompt.js';
 import type { TextBuffer } from './shared/text-buffer.js';
 import { Config } from '@google/gemini-cli-core';
@@ -1226,11 +1227,12 @@ describe('InputPrompt', () => {
       stdin.write('\x12');
       await wait();
       stdin.write('\x1B');
-      await wait();
 
-      const frame = stdout.lastFrame();
-      expect(frame).not.toContain('(r:)');
-      expect(frame).not.toContain('echo hello');
+      await waitFor(() => {
+        expect(stdout.lastFrame()).not.toContain('(r:)');
+      });
+
+      expect(stdout.lastFrame()).not.toContain('echo hello');
 
       unmount();
     });
@@ -1240,9 +1242,11 @@ describe('InputPrompt', () => {
       stdin.write('\x12');
       await wait();
       stdin.write('\t');
-      await wait();
 
-      expect(stdout.lastFrame()).not.toContain('(r:)');
+      await waitFor(() => {
+        expect(stdout.lastFrame()).not.toContain('(r:)');
+      });
+
       expect(props.buffer.setText).toHaveBeenCalledWith('echo hello');
       unmount();
     });
@@ -1253,9 +1257,11 @@ describe('InputPrompt', () => {
       await wait();
       expect(stdout.lastFrame()).toContain('(r:)');
       stdin.write('\r');
-      await wait();
 
-      expect(stdout.lastFrame()).not.toContain('(r:)');
+      await waitFor(() => {
+        expect(stdout.lastFrame()).not.toContain('(r:)');
+      });
+
       expect(props.onSubmit).toHaveBeenCalledWith('echo hello');
       unmount();
     });
@@ -1268,9 +1274,10 @@ describe('InputPrompt', () => {
       await wait();
       expect(stdout.lastFrame()).toContain('(r:)');
       stdin.write('\x1B');
-      await wait();
 
-      expect(stdout.lastFrame()).not.toContain('(r:)');
+      await waitFor(() => {
+        expect(stdout.lastFrame()).not.toContain('(r:)');
+      });
       expect(props.buffer.text).toBe('initial text');
       expect(props.buffer.cursor).toEqual([0, 3]);