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();
     });