From a5c81e3fe07d264dd467d9d2aebb8ae0bc32bfb3 Mon Sep 17 00:00:00 2001
From: Jacob Richman <jacob314@gmail.com>
Date: Thu, 14 Aug 2025 16:39:12 -0700
Subject: [PATCH] Fix flake due to using wait instead of waitFor (#6277)

---
 packages/cli/src/ui/components/InputPrompt.test.tsx | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/packages/cli/src/ui/components/InputPrompt.test.tsx b/packages/cli/src/ui/components/InputPrompt.test.tsx
index a29a095a..ec2d7441 100644
--- a/packages/cli/src/ui/components/InputPrompt.test.tsx
+++ b/packages/cli/src/ui/components/InputPrompt.test.tsx
@@ -1217,17 +1217,18 @@ describe('InputPrompt', () => {
       props.buffer.setText('some text');
 
       const { stdin, unmount } = render(<InputPrompt {...props} />);
-      await wait();
 
       stdin.write('\x1B');
-      await wait();
 
-      expect(onEscapePromptChange).toHaveBeenCalledWith(true);
+      await waitFor(() => {
+        expect(onEscapePromptChange).toHaveBeenCalledWith(true);
+      });
 
       stdin.write('a');
-      await wait();
 
-      expect(onEscapePromptChange).toHaveBeenCalledWith(false);
+      await waitFor(() => {
+        expect(onEscapePromptChange).toHaveBeenCalledWith(false);
+      });
       unmount();
     });