1 files changed, 45 insertions, 0 deletions
diff --git a/packages/cli/src/ui/components/InputPrompt.test.tsx b/packages/cli/src/ui/components/InputPrompt.test.tsx
index a1894002..60ba648d 100644
--- a/packages/cli/src/ui/components/InputPrompt.test.tsx
+++ b/packages/cli/src/ui/components/InputPrompt.test.tsx
@@ -171,6 +171,7 @@ describe('InputPrompt', () => {
       config: {
         getProjectRoot: () => path.join('test', 'project'),
         getTargetDir: () => path.join('test', 'project', 'src'),
+        getVimMode: () => false,
       } as unknown as Config,
       slashCommands: mockSlashCommands,
       commandContext: mockCommandContext,
@@ -1076,4 +1077,48 @@ describe('InputPrompt', () => {
       unmount();
     });
   });
+
+  describe('vim mode', () => {
+    it('should not call buffer.handleInput when vim mode is enabled and vim handles the input', async () => {
+      props.vimModeEnabled = true;
+      props.vimHandleInput = vi.fn().mockReturnValue(true); // Mock that vim handled it.
+      const { stdin, unmount } = render(<InputPrompt {...props} />);
+      await wait();
+
+      stdin.write('i');
+      await wait();
+
+      expect(props.vimHandleInput).toHaveBeenCalled();
+      expect(mockBuffer.handleInput).not.toHaveBeenCalled();
+      unmount();
+    });
+
+    it('should call buffer.handleInput when vim mode is enabled but vim does not handle the input', async () => {
+      props.vimModeEnabled = true;
+      props.vimHandleInput = vi.fn().mockReturnValue(false); // Mock that vim did NOT handle it.
+      const { stdin, unmount } = render(<InputPrompt {...props} />);
+      await wait();
+
+      stdin.write('i');
+      await wait();
+
+      expect(props.vimHandleInput).toHaveBeenCalled();
+      expect(mockBuffer.handleInput).toHaveBeenCalled();
+      unmount();
+    });
+
+    it('should call handleInput when vim mode is disabled', async () => {
+      // Mock vimHandleInput to return false (vim didn't handle the input)
+      props.vimHandleInput = vi.fn().mockReturnValue(false);
+      const { stdin, unmount } = render(<InputPrompt {...props} />);
+      await wait();
+
+      stdin.write('i');
+      await wait();
+
+      expect(props.vimHandleInput).toHaveBeenCalled();
+      expect(mockBuffer.handleInput).toHaveBeenCalled();
+      unmount();
+    });
+  });
 });