aider/tests/basic/test_io.py

import os
import unittest
from pathlib import Path
from unittest.mock import MagicMock, patch

from prompt_toolkit.completion import CompleteEvent
from prompt_toolkit.document import Document
from rich.text import Text

from aider.dump import dump  # noqa: F401
from aider.io import AutoCompleter, ConfirmGroup, InputOutput
from aider.utils import ChdirTemporaryDirectory


class TestInputOutput(unittest.TestCase):
    def test_line_endings_validation(self):
        # Test valid line endings
        for ending in ["platform", "lf", "crlf"]:
            io = InputOutput(line_endings=ending)
            self.assertEqual(
                io.newline, None if ending == "platform" else "\n" if ending == "lf" else "\r\n"
            )

        # Test invalid line endings
        with self.assertRaises(ValueError) as cm:
            io = InputOutput(line_endings="invalid")
        self.assertIn("Invalid line_endings value: invalid", str(cm.exception))
        # Check each valid option is in the error message
        self.assertIn("platform", str(cm.exception))
        self.assertIn("crlf", str(cm.exception))
        self.assertIn("lf", str(cm.exception))

    def test_no_color_environment_variable(self):
        with patch.dict(os.environ, {"NO_COLOR": "1"}):
            io = InputOutput(fancy_input=False)
            self.assertFalse(io.pretty)

    def test_color_initialization(self):
        """Test that color values are properly initialized with # prefix"""
        # Test with hex colors without #
        io = InputOutput(
            user_input_color="00cc00",
            tool_error_color="FF2222",
            tool_warning_color="FFA500",
            assistant_output_color="0088ff",
            pretty=True,
        )

        # Check that # was added to hex colors
        self.assertEqual(io.user_input_color, "#00cc00")
        self.assertEqual(io.tool_error_color, "#FF2222")
        self.assertEqual(io.tool_warning_color, "#FFA500")  # Already had #
        self.assertEqual(io.assistant_output_color, "#0088ff")

        # Test with named colors (should be unchanged)
        io = InputOutput(user_input_color="blue", tool_error_color="red", pretty=True)

        self.assertEqual(io.user_input_color, "blue")
        self.assertEqual(io.tool_error_color, "red")

        # Test with pretty=False (should not modify colors)
        io = InputOutput(user_input_color="00cc00", tool_error_color="FF2222", pretty=False)

        self.assertIsNone(io.user_input_color)
        self.assertIsNone(io.tool_error_color)

    def test_dumb_terminal(self):
        with patch.dict(os.environ, {"TERM": "dumb"}):
            io = InputOutput(fancy_input=True)
            self.assertTrue(io.is_dumb_terminal)
            self.assertFalse(io.pretty)
            self.assertIsNone(io.prompt_session)

    def test_autocompleter_get_command_completions(self):
        # Step 3: Mock the commands object
        commands = MagicMock()
        commands.get_commands.return_value = ["/help", "/add", "/drop"]
        commands.matching_commands.side_effect = lambda inp: (
            [cmd for cmd in commands.get_commands() if cmd.startswith(inp.strip().split()[0])],
            inp.strip().split()[0],
            " ".join(inp.strip().split()[1:]),
        )
        commands.get_raw_completions.return_value = None
        commands.get_completions.side_effect = lambda cmd: (
            ["file1.txt", "file2.txt"] if cmd == "/add" else None
        )

        # Step 4: Create an instance of AutoCompleter
        root = ""
        rel_fnames = []
        addable_rel_fnames = []
        autocompleter = AutoCompleter(
            root=root,
            rel_fnames=rel_fnames,
            addable_rel_fnames=addable_rel_fnames,
            commands=commands,
            encoding="utf-8",
        )

        # Step 5: Set up test cases
        test_cases = [
            # Input text, Expected completion texts
            ("/", ["/help", "/add", "/drop"]),
            ("/a", ["/add"]),
            ("/add f", ["file1.txt", "file2.txt"]),
        ]

        # Step 6: Iterate through test cases
        for text, expected_completions in test_cases:
            document = Document(text=text)
            complete_event = CompleteEvent()
            words = text.strip().split()

            # Call get_command_completions
            completions = list(
                autocompleter.get_command_completions(
                    document,
                    complete_event,
                    text,
                    words,
                )
            )

            # Extract completion texts
            completion_texts = [comp.text for comp in completions]

            # Assert that the completions match expected results
            self.assertEqual(set(completion_texts), set(expected_completions))

    def test_autocompleter_with_non_existent_file(self):
        root = ""
        rel_fnames = ["non_existent_file.txt"]
        addable_rel_fnames = []
        commands = None
        autocompleter = AutoCompleter(root, rel_fnames, addable_rel_fnames, commands, "utf-8")
        self.assertEqual(autocompleter.words, set(rel_fnames))

    def test_autocompleter_with_unicode_file(self):
        with ChdirTemporaryDirectory():
            root = ""
            fname = "file.py"
            rel_fnames = [fname]
            addable_rel_fnames = []
            commands = None
            autocompleter = AutoCompleter(root, rel_fnames, addable_rel_fnames, commands, "utf-8")
            self.assertEqual(autocompleter.words, set(rel_fnames))

            Path(fname).write_text("def hello(): pass\n")
            autocompleter = AutoCompleter(root, rel_fnames, addable_rel_fnames, commands, "utf-8")
            autocompleter.tokenize()
            dump(autocompleter.words)
            self.assertEqual(autocompleter.words, set(rel_fnames + [("hello", "`hello`")]))

            encoding = "utf-16"
            some_content_which_will_error_if_read_with_encoding_utf8 = "ÅÍÎÏ".encode(encoding)
            with open(fname, "wb") as f:
                f.write(some_content_which_will_error_if_read_with_encoding_utf8)

            autocompleter = AutoCompleter(root, rel_fnames, addable_rel_fnames, commands, "utf-8")
            self.assertEqual(autocompleter.words, set(rel_fnames))

    @patch("builtins.input", return_value="test input")
    def test_get_input_is_a_directory_error(self, mock_input):
        io = InputOutput(pretty=False, fancy_input=False)  # Windows tests throw UnicodeDecodeError
        root = "/"
        rel_fnames = ["existing_file.txt"]
        addable_rel_fnames = ["new_file.txt"]
        commands = MagicMock()

        # Simulate IsADirectoryError
        with patch("aider.io.open", side_effect=IsADirectoryError):
            result = io.get_input(root, rel_fnames, addable_rel_fnames, commands)
            self.assertEqual(result, "test input")
            mock_input.assert_called_once()

    @patch("builtins.input")
    def test_confirm_ask_explicit_yes_required(self, mock_input):
        io = InputOutput(pretty=False, fancy_input=False)

        # Test case 1: explicit_yes_required=True, self.yes=True
        io.yes = True
        result = io.confirm_ask("Are you sure?", explicit_yes_required=True)
        self.assertFalse(result)
        mock_input.assert_not_called()

        # Test case 2: explicit_yes_required=True, self.yes=False
        io.yes = False
        result = io.confirm_ask("Are you sure?", explicit_yes_required=True)
        self.assertFalse(result)
        mock_input.assert_not_called()

        # Test case 3: explicit_yes_required=True, user input required
        io.yes = None
        mock_input.return_value = "y"
        result = io.confirm_ask("Are you sure?", explicit_yes_required=True)
        self.assertTrue(result)
        mock_input.assert_called_once()

        # Reset mock_input
        mock_input.reset_mock()

        # Test case 4: explicit_yes_required=False, self.yes=True
        io.yes = True
        result = io.confirm_ask("Are you sure?", explicit_yes_required=False)
        self.assertTrue(result)
        mock_input.assert_not_called()

    @patch("builtins.input")
    def test_confirm_ask_with_group(self, mock_input):
        io = InputOutput(pretty=False, fancy_input=False)
        group = ConfirmGroup()

        # Test case 1: No group preference, user selects 'All'
        mock_input.return_value = "a"
        result = io.confirm_ask("Are you sure?", group=group)
        self.assertTrue(result)
        self.assertEqual(group.preference, "all")
        mock_input.assert_called_once()
        mock_input.reset_mock()

        # Test case 2: Group preference is 'All', should not prompt
        result = io.confirm_ask("Are you sure?", group=group)
        self.assertTrue(result)
        mock_input.assert_not_called()

        # Test case 3: No group preference, user selects 'Skip all'
        group.preference = None
        mock_input.return_value = "s"
        result = io.confirm_ask("Are you sure?", group=group)
        self.assertFalse(result)
        self.assertEqual(group.preference, "skip")
        mock_input.assert_called_once()
        mock_input.reset_mock()

        # Test case 4: Group preference is 'Skip all', should not prompt
        result = io.confirm_ask("Are you sure?", group=group)
        self.assertFalse(result)
        mock_input.assert_not_called()

        # Test case 5: explicit_yes_required=True, should not offer 'All' option
        group.preference = None
        mock_input.return_value = "y"
        result = io.confirm_ask("Are you sure?", group=group, explicit_yes_required=True)
        self.assertTrue(result)
        self.assertIsNone(group.preference)
        mock_input.assert_called_once()
        self.assertNotIn("(A)ll", mock_input.call_args[0][0])
        mock_input.reset_mock()

    @patch("builtins.input")
    def test_confirm_ask_yes_no(self, mock_input):
        io = InputOutput(pretty=False, fancy_input=False)

        # Test case 1: User selects 'Yes'
        mock_input.return_value = "y"
        result = io.confirm_ask("Are you sure?")
        self.assertTrue(result)
        mock_input.assert_called_once()
        mock_input.reset_mock()

        # Test case 2: User selects 'No'
        mock_input.return_value = "n"
        result = io.confirm_ask("Are you sure?")
        self.assertFalse(result)
        mock_input.assert_called_once()
        mock_input.reset_mock()

        # Test case 3: Empty input (default to Yes)
        mock_input.return_value = ""
        result = io.confirm_ask("Are you sure?")
        self.assertTrue(result)
        mock_input.assert_called_once()
        mock_input.reset_mock()

        # Test case 4: 'skip' functions as 'no' without group
        mock_input.return_value = "s"
        result = io.confirm_ask("Are you sure?")
        self.assertFalse(result)
        mock_input.assert_called_once()
        mock_input.reset_mock()

        # Test case 5: 'all' functions as 'yes' without group
        mock_input.return_value = "a"
        result = io.confirm_ask("Are you sure?")
        self.assertTrue(result)
        mock_input.assert_called_once()
        mock_input.reset_mock()

        # Test case 6: Full word 'skip' functions as 'no' without group
        mock_input.return_value = "skip"
        result = io.confirm_ask("Are you sure?")
        self.assertFalse(result)
        mock_input.assert_called_once()
        mock_input.reset_mock()

        # Test case 7: Full word 'all' functions as 'yes' without group
        mock_input.return_value = "all"
        result = io.confirm_ask("Are you sure?")
        self.assertTrue(result)
        mock_input.assert_called_once()
        mock_input.reset_mock()

    @patch("builtins.input", side_effect=["d"])
    def test_confirm_ask_allow_never(self, mock_input):
        """Test the 'don't ask again' functionality in confirm_ask"""
        io = InputOutput(pretty=False, fancy_input=False)

        # First call: user selects "Don't ask again"
        result = io.confirm_ask("Are you sure?", allow_never=True)
        self.assertFalse(result)
        mock_input.assert_called_once()
        self.assertIn(("Are you sure?", None), io.never_prompts)

        # Reset the mock to check for further calls
        mock_input.reset_mock()

        # Second call: should not prompt, immediately return False
        result = io.confirm_ask("Are you sure?", allow_never=True)
        self.assertFalse(result)
        mock_input.assert_not_called()

        # Test with subject parameter
        mock_input.reset_mock()
        mock_input.side_effect = ["d"]
        result = io.confirm_ask("Confirm action?", subject="Subject Text", allow_never=True)
        self.assertFalse(result)
        mock_input.assert_called_once()
        self.assertIn(("Confirm action?", "Subject Text"), io.never_prompts)

        # Subsequent call with the same question and subject
        mock_input.reset_mock()
        result = io.confirm_ask("Confirm action?", subject="Subject Text", allow_never=True)
        self.assertFalse(result)
        mock_input.assert_not_called()

        # Test that allow_never=False does not add to never_prompts
        mock_input.reset_mock()
        mock_input.side_effect = ["d", "n"]
        result = io.confirm_ask("Do you want to proceed?", allow_never=False)
        self.assertFalse(result)
        self.assertEqual(mock_input.call_count, 2)
        self.assertNotIn(("Do you want to proceed?", None), io.never_prompts)


class TestInputOutputMultilineMode(unittest.TestCase):
    def setUp(self):
        self.io = InputOutput(fancy_input=True)
        self.io.prompt_session = MagicMock()

    def test_toggle_multiline_mode(self):
        """Test that toggling multiline mode works correctly"""
        # Start in single-line mode
        self.io.multiline_mode = False

        # Toggle to multiline mode
        self.io.toggle_multiline_mode()
        self.assertTrue(self.io.multiline_mode)

        # Toggle back to single-line mode
        self.io.toggle_multiline_mode()
        self.assertFalse(self.io.multiline_mode)

    def test_tool_message_unicode_fallback(self):
        """Test that Unicode messages are properly converted to ASCII with replacement"""
        io = InputOutput(pretty=False, fancy_input=False)

        # Create a message with invalid Unicode that can't be encoded in UTF-8
        # Using a surrogate pair that's invalid in UTF-8
        invalid_unicode = "Hello \ud800World"

        # Mock console.print to capture the output
        with patch.object(io.console, "print") as mock_print:
            # First call will raise UnicodeEncodeError
            mock_print.side_effect = [UnicodeEncodeError("utf-8", "", 0, 1, "invalid"), None]

            io._tool_message(invalid_unicode)

            # Verify that the message was converted to ASCII with replacement
            self.assertEqual(mock_print.call_count, 2)
            args, kwargs = mock_print.call_args
            converted_message = args[0]

            # The invalid Unicode should be replaced with '?'
            self.assertEqual(converted_message, "Hello ?World")

    def test_multiline_mode_restored_after_interrupt(self):
        """Test that multiline mode is restored after KeyboardInterrupt"""
        io = InputOutput(fancy_input=True)
        io.prompt_session = MagicMock()

        # Start in multiline mode
        io.multiline_mode = True

        # Mock prompt() to raise KeyboardInterrupt
        io.prompt_session.prompt.side_effect = KeyboardInterrupt

        # Test confirm_ask()
        with self.assertRaises(KeyboardInterrupt):
            io.confirm_ask("Test question?")
        self.assertTrue(io.multiline_mode)  # Should be restored

        # Test prompt_ask()
        with self.assertRaises(KeyboardInterrupt):
            io.prompt_ask("Test prompt?")
        self.assertTrue(io.multiline_mode)  # Should be restored

    def test_multiline_mode_restored_after_normal_exit(self):
        """Test that multiline mode is restored after normal exit"""
        io = InputOutput(fancy_input=True)
        io.prompt_session = MagicMock()

        # Start in multiline mode
        io.multiline_mode = True

        # Mock prompt() to return normally
        io.prompt_session.prompt.return_value = "y"

        # Test confirm_ask()
        io.confirm_ask("Test question?")
        self.assertTrue(io.multiline_mode)  # Should be restored

        # Test prompt_ask()
        io.prompt_ask("Test prompt?")
        self.assertTrue(io.multiline_mode)  # Should be restored

    def test_ensure_hash_prefix(self):
        """Test that ensure_hash_prefix correctly adds # to valid hex colors"""
        from aider.io import ensure_hash_prefix

        # Test valid hex colors without #
        self.assertEqual(ensure_hash_prefix("000"), "#000")
        self.assertEqual(ensure_hash_prefix("fff"), "#fff")
        self.assertEqual(ensure_hash_prefix("F00"), "#F00")
        self.assertEqual(ensure_hash_prefix("123456"), "#123456")
        self.assertEqual(ensure_hash_prefix("abcdef"), "#abcdef")
        self.assertEqual(ensure_hash_prefix("ABCDEF"), "#ABCDEF")

        # Test hex colors that already have #
        self.assertEqual(ensure_hash_prefix("#000"), "#000")
        self.assertEqual(ensure_hash_prefix("#123456"), "#123456")

        # Test invalid inputs (should return unchanged)
        self.assertEqual(ensure_hash_prefix(""), "")
        self.assertEqual(ensure_hash_prefix(None), None)
        self.assertEqual(ensure_hash_prefix("red"), "red")  # Named color
        self.assertEqual(ensure_hash_prefix("12345"), "12345")  # Wrong length
        self.assertEqual(ensure_hash_prefix("1234567"), "1234567")  # Wrong length
        self.assertEqual(ensure_hash_prefix("xyz"), "xyz")  # Invalid hex chars
        self.assertEqual(ensure_hash_prefix("12345g"), "12345g")  # Invalid hex chars

    def test_tool_output_color_handling(self):
        """Test that tool_output correctly handles hex colors without # prefix"""
        from unittest.mock import patch

        # Create IO with hex color without # for tool_output_color
        io = InputOutput(tool_output_color="FFA500", pretty=True)

        # Patch console.print to avoid actual printing
        with patch.object(io.console, "print") as mock_print:
            # This would raise ColorParseError without the fix
            io.tool_output("Test message")

            # Verify the call was made without error
            mock_print.assert_called_once()

            # Verify the style was correctly created with # prefix
            # The first argument is the message, second would be the style
            kwargs = mock_print.call_args.kwargs
            self.assertIn("style", kwargs)

        # Test with other hex color
        io = InputOutput(tool_output_color="00FF00", pretty=True)
        with patch.object(io.console, "print") as mock_print:
            io.tool_output("Test message")
            mock_print.assert_called_once()


@patch("aider.io.is_dumb_terminal", return_value=False)
@patch.dict(os.environ, {"NO_COLOR": ""})
class TestInputOutputFormatFiles(unittest.TestCase):
    def test_format_files_for_input_pretty_false(self, mock_is_dumb_terminal):
        io = InputOutput(pretty=False, fancy_input=False)
        rel_fnames = ["file1.txt", "file[markup].txt", "ro_file.txt"]
        rel_read_only_fnames = ["ro_file.txt"]

        expected_output = "file1.txt\nfile[markup].txt\nro_file.txt (read only)\n"
        # Sort the expected lines because the order of editable vs read-only might vary
        # depending on internal sorting, but the content should be the same.
        # The method sorts editable_files and read_only_files separately.
        # The final output joins sorted(read_only_files) + sorted(editable_files)

        # Based on current implementation:
        # read_only_files = ["ro_file.txt (read only)"]
        # editable_files = ["file1.txt", "file[markup].txt"]
        # output = "\n".join(read_only_files + editable_files) + "\n"

        # Correct expected output based on implementation:
        expected_output_lines = sorted(
            [
                "ro_file.txt (read only)",
                "file1.txt",
                "file[markup].txt",
            ]
        )
        expected_output = "\n".join(expected_output_lines) + "\n"

        actual_output = io.format_files_for_input(rel_fnames, rel_read_only_fnames)

        # Normalizing actual output by splitting, sorting, and rejoining
        actual_output_lines = sorted(filter(None, actual_output.splitlines()))
        normalized_actual_output = "\n".join(actual_output_lines) + "\n"

        self.assertEqual(normalized_actual_output, expected_output)

    @patch("aider.io.Columns")
    @patch("os.path.abspath")
    @patch("os.path.join")
    def test_format_files_for_input_pretty_true_no_files(
        self, mock_join, mock_abspath, mock_columns, mock_is_dumb_terminal
    ):
        io = InputOutput(pretty=True, root="test_root")
        io.format_files_for_input([], [])
        mock_columns.assert_not_called()

    @patch("aider.io.Columns")
    @patch("os.path.abspath")
    @patch("os.path.join")
    def test_format_files_for_input_pretty_true_editable_only(
        self, mock_join, mock_abspath, mock_columns, mock_is_dumb_terminal
    ):
        io = InputOutput(pretty=True, root="test_root")
        rel_fnames = ["edit1.txt", "edit[markup].txt"]

        io.format_files_for_input(rel_fnames, [])

        mock_columns.assert_called_once()
        args, _ = mock_columns.call_args
        renderables = args[0]

        self.assertEqual(len(renderables), 2)
        self.assertIsInstance(renderables[0], Text)
        self.assertEqual(renderables[0].plain, "edit1.txt")
        self.assertIsInstance(renderables[1], Text)
        self.assertEqual(renderables[1].plain, "edit[markup].txt")

    @patch("aider.io.Columns")
    @patch("os.path.abspath")
    @patch("os.path.join")
    def test_format_files_for_input_pretty_true_readonly_only(
        self, mock_join, mock_abspath, mock_columns, mock_is_dumb_terminal
    ):
        io = InputOutput(pretty=True, root="test_root")

        # Mock path functions to ensure rel_path is chosen by the shortener logic
        mock_join.side_effect = lambda *args: "/".join(args)
        mock_abspath.side_effect = lambda p: "/ABS_PREFIX_VERY_LONG/" + os.path.normpath(p)

        rel_read_only_fnames = ["ro1.txt", "ro[markup].txt"]
        # When all files in chat are read-only
        rel_fnames = list(rel_read_only_fnames)

        io.format_files_for_input(rel_fnames, rel_read_only_fnames)

        self.assertEqual(mock_columns.call_count, 2)
        args, _ = mock_columns.call_args
        renderables = args[0]

        self.assertEqual(len(renderables), 3)  # Readonly: + 2 files
        self.assertIsInstance(renderables[0], Text)
        self.assertEqual(renderables[0].plain, "Readonly:")
        self.assertIsInstance(renderables[1], Text)
        self.assertEqual(renderables[1].plain, "ro1.txt")
        self.assertIsInstance(renderables[2], Text)
        self.assertEqual(renderables[2].plain, "ro[markup].txt")

    @patch("aider.io.Columns")
    @patch("os.path.abspath")
    @patch("os.path.join")
    def test_format_files_for_input_pretty_true_mixed_files(
        self, mock_join, mock_abspath, mock_columns, mock_is_dumb_terminal
    ):
        io = InputOutput(pretty=True, root="test_root")

        mock_join.side_effect = lambda *args: "/".join(args)
        mock_abspath.side_effect = lambda p: "/ABS_PREFIX_VERY_LONG/" + os.path.normpath(p)

        rel_fnames = ["edit1.txt", "edit[markup].txt", "ro1.txt", "ro[markup].txt"]
        rel_read_only_fnames = ["ro1.txt", "ro[markup].txt"]

        io.format_files_for_input(rel_fnames, rel_read_only_fnames)

        self.assertEqual(mock_columns.call_count, 4)

        # Check arguments for the first rendering of read-only files (call 0)
        args_ro, _ = mock_columns.call_args_list[0]
        renderables_ro = args_ro[0]
        self.assertEqual(
            renderables_ro, [Text("Readonly:"), Text("ro1.txt"), Text("ro[markup].txt")]
        )

        # Check arguments for the first rendering of editable files (call 2)
        args_ed, _ = mock_columns.call_args_list[2]
        renderables_ed = args_ed[0]
        self.assertEqual(
            renderables_ed, [Text("Editable:"), Text("edit1.txt"), Text("edit[markup].txt")]
        )


if __name__ == "__main__":
    unittest.main()