diff --git a/tests/test_editblock.py b/tests/test_editblock.py index bed2a5f18..6aea04124 100644 --- a/tests/test_editblock.py +++ b/tests/test_editblock.py @@ -1,11 +1,26 @@ # flake8: noqa: E501 +import tempfile import unittest +from pathlib import Path +from unittest.mock import MagicMock, patch +from aider import models +from aider.coders import Coder from aider.coders import editblock_coder as eb +from aider.dump import dump # noqa: F401 +from aider.io import InputOutput class TestUtils(unittest.TestCase): + def setUp(self): + self.patcher = patch("aider.coders.base_coder.check_model_availability") + self.mock_check = self.patcher.start() + self.mock_check.return_value = True + + def tearDown(self): + self.patcher.stop() + def test_replace_most_similar_chunk(self): whole = "This is a sample text.\nAnother line of text.\nYet another line.\n" part = "This is a sample text" @@ -223,6 +238,42 @@ These changes replace the `subprocess.run` patches with `subprocess.check_output result = eb.replace_part_with_missing_leading_whitespace(whole, part, replace) self.assertEqual(result, expected_output) + def test_choose_fence(self): + # Create a few temporary files + _, file1 = tempfile.mkstemp() + + with open(file1, "w", encoding="utf-8") as f: + f.write("one\ntwo\nthree\n") + + files = [file1] + + # Initialize the Coder object with the mocked IO and mocked repo + coder = Coder.create( + models.GPT4, "diff", io=InputOutput(), openai_api_key="fake_key", fnames=files + ) + + def mock_send(*args, **kwargs): + coder.partial_response_content = f""" +Do this: + +{Path(file1).name} +<<<<<<< ORIGINAL +two +======= +new +>>>>>>> UPDATED + +""" + coder.partial_response_function_call = dict() + + coder.send = MagicMock(side_effect=mock_send) + + # Call the run method with a message + coder.run(with_message="hi") + + content = open(file1, encoding="utf-8").read() + self.assertEqual(content, "one\nnew\nthree\n") + if __name__ == "__main__": unittest.main()