fix: Mock Coder.send() instead of Coder.create and Model

This commit is contained in:
Paul Gauthier (aider) 2024-08-01 15:37:34 -03:00
parent b29a6e8fa4
commit 4eb8be9d06

View file

@ -7,34 +7,30 @@ from aider.models import Model
class TestScriptingAPI(unittest.TestCase): class TestScriptingAPI(unittest.TestCase):
@patch("aider.coders.Coder.create") @patch("aider.coders.base_coder.Coder.send")
@patch("aider.models.Model") def test_basic_scripting(self, mock_send):
def test_basic_scripting(self, mock_model, mock_coder_create):
# Setup # Setup
mock_coder = MagicMock() mock_send.return_value = "Changes applied successfully."
mock_coder_create.return_value = mock_coder
# Test script # Test script
fnames = ["greeting.py"] fnames = ["greeting.py"]
model = Model("gpt-4-turbo") model = Model("gpt-4-turbo")
coder = Coder.create(main_model=model, fnames=fnames) coder = Coder.create(main_model=model, fnames=fnames)
coder.run("make a script that prints hello world") result1 = coder.run("make a script that prints hello world")
coder.run("make it say goodbye") result2 = coder.run("make it say goodbye")
# Assertions # Assertions
mock_model.assert_called_once_with("gpt-4-turbo") self.assertEqual(mock_send.call_count, 2)
mock_coder_create.assert_called_once_with(main_model=model, fnames=fnames) mock_send.assert_any_call([{'role': 'user', 'content': 'make a script that prints hello world'}])
self.assertEqual(mock_coder.run.call_count, 2) mock_send.assert_any_call([{'role': 'user', 'content': 'make it say goodbye'}])
mock_coder.run.assert_any_call("make a script that prints hello world") self.assertEqual(result1, "Changes applied successfully.")
mock_coder.run.assert_any_call("make it say goodbye") self.assertEqual(result2, "Changes applied successfully.")
@patch("aider.coders.Coder.create") @patch("aider.coders.base_coder.Coder.send")
@patch("aider.models.Model") def test_scripting_with_io(self, mock_send):
def test_scripting_with_io(self, mock_model, mock_coder_create):
# Setup # Setup
mock_coder = MagicMock() mock_send.return_value = "New function added successfully."
mock_coder_create.return_value = mock_coder
# Test script # Test script
fnames = ["greeting.py"] fnames = ["greeting.py"]
@ -42,12 +38,11 @@ class TestScriptingAPI(unittest.TestCase):
io = InputOutput(yes=True) io = InputOutput(yes=True)
coder = Coder.create(main_model=model, fnames=fnames, io=io) coder = Coder.create(main_model=model, fnames=fnames, io=io)
coder.run("add a new function") result = coder.run("add a new function")
# Assertions # Assertions
mock_model.assert_called_once_with("gpt-4-turbo") mock_send.assert_called_once_with([{'role': 'user', 'content': 'add a new function'}])
mock_coder_create.assert_called_once_with(main_model=model, fnames=fnames, io=io) self.assertEqual(result, "New function added successfully.")
mock_coder.run.assert_called_once_with("add a new function")
self.assertTrue(io.yes) # Check that 'yes' is set to True self.assertTrue(io.yes) # Check that 'yes' is set to True