tests/basic/test_wholefile.py
Rama main · réplica de lectura del repositorio original.
import os
import shutil
import tempfile
import unittest
from pathlib import Path
from unittest.mock import MagicMock
from aider.coders import Coder
from aider.coders.wholefile_coder import WholeFileCoder
from aider.dump import dump # noqa: F401
from aider.io import InputOutput
from aider.models import Model
class TestWholeFileCoder(unittest.TestCase):
def setUp(self):
self.original_cwd = os.getcwd()
self.tempdir = tempfile.mkdtemp()
os.chdir(self.tempdir)
self.GPT35 = Model("gpt-3.5-turbo")
def tearDown(self):
os.chdir(self.original_cwd)
shutil.rmtree(self.tempdir, ignore_errors=True)
def test_no_files(self):
# Initialize WholeFileCoder with the temporary directory
io = InputOutput(yes=True)
coder = WholeFileCoder(main_model=self.GPT35, io=io, fnames=[])
coder.partial_response_content = (
'To print "Hello, World!" in most programming languages, you can use the following'
' code:\n\n```python\nprint("Hello, World!")\n```\n\nThis code will output "Hello,'
' World!" to the console.'
)
# This is throwing ValueError!
coder.render_incremental_response(True)
def test_no_files_new_file_should_ask(self):
io = InputOutput(yes=False) # <- yes=FALSE
coder = WholeFileCoder(main_model=self.GPT35, io=io, fnames=[])
coder.partial_response_content = (
'To print "Hello, World!" in most programming languages, you can use the following'
' code:\n\nfoo.js\n```python\nprint("Hello, World!")\n```\n\nThis code will output'
' "Hello, World!" to the console.'
)
coder.apply_updates()
self.assertFalse(Path("foo.js").exists())
def test_update_files(self):
# Create a sample file in the temporary directory
sample_file = "sample.txt"
with open(sample_file, "w") as f:
f.write("Original content\n")
# Initialize WholeFileCoder with the temporary directory
io = InputOutput(yes=True)
coder = WholeFileCoder(main_model=self.GPT35, io=io, fnames=[sample_file])
# Set the partial response content with the updated content
coder.partial_response_content = f"{sample_file}\n```\nUpdated content\n```"
# Call update_files method
edited_files = coder.apply_updates()
# Check if the sample file was updated
self.assertIn("sample.txt", edited_files)
# Check if the content of the sample file was updated
with open(sample_file, "r") as f:
updated_content = f.read()
self.assertEqual(updated_content, "Updated content\n")
def test_update_files_live_diff(self):
# Create a sample file in the temporary directory
sample_file = "sample.txt"
with open(sample_file, "w") as f:
f.write("\n".join(map(str, range(0, 100))))
# Initialize WholeFileCoder with the temporary directory
io = InputOutput(yes=True)
coder = WholeFileCoder(main_model=self.GPT35, io=io, fnames=[sample_file])
# Set the partial response content with the updated content
coder.partial_response_content = f"{sample_file}\n```\n0\n\1\n2\n"
lines = coder.get_edits(mode="diff").splitlines()
# the live diff should be concise, since we haven't changed anything yet
self.assertLess(len(lines), 20)
def test_update_files_with_existing_fence(self):
# Create a sample file in the temporary directory
sample_file = "sample.txt"
original_content = """
Here is some quoted text:
```
Quote!
```
"""
with open(sample_file, "w") as f:
f.write(original_content)
# Initialize WholeFileCoder with the temporary directory
io = InputOutput(yes=True)
coder = WholeFileCoder(main_model=self.GPT35, io=io, fnames=[sample_file])
coder.choose_fence()
self.assertNotEqual(coder.fence[0], "```")
# Set the partial response content with the updated content
coder.partial_response_content = (
f"{sample_file}\n{coder.fence[0]}\nUpdated content\n{coder.fence[1]}"
)
# Call update_files method
edited_files = coder.apply_updates()
# Check if the sample file was updated
self.assertIn("sample.txt", edited_files)
# Check if the content of the sample file was updated
with open(sample_file, "r") as f:
updated_content = f.read()
self.assertEqual(updated_content, "Updated content\n")
def test_update_files_bogus_path_prefix(self):
# Create a sample file in the temporary directory
sample_file = "sample.txt"
with open(sample_file, "w") as f:
f.write("Original content\n")
# Initialize WholeFileCoder with the temporary directory
io = InputOutput(yes=True)
coder = WholeFileCoder(main_model=self.GPT35, io=io, fnames=[sample_file])
# Set the partial response content with the updated content
# With path/to/ prepended onto the filename
coder.partial_response_content = f"path/to/{sample_file}\n```\nUpdated content\n```"
# Call update_files method
edited_files = coder.apply_updates()
# Check if the sample file was updated
self.assertIn("sample.txt", edited_files)
# Check if the content of the sample file was updated
with open(sample_file, "r") as f:
updated_content = f.read()
self.assertEqual(updated_content, "Updated content\n")
def test_update_files_not_in_chat(self):
# Create a sample file in the temporary directory
sample_file = "sample.txt"
with open(sample_file, "w") as f:
f.write("Original content\n")
# Initialize WholeFileCoder with the temporary directory
io = InputOutput(yes=True)
coder = WholeFileCoder(main_model=self.GPT35, io=io)
# Set the partial response content with the updated content
coder.partial_response_content = f"{sample_file}\n```\nUpdated content\n```"
# Call update_files method
edited_files = coder.apply_updates()
# Check if the sample file was updated
self.assertIn("sample.txt", edited_files)
# Check if the content of the sample file was updated
with open(sample_file, "r") as f:
updated_content = f.read()
self.assertEqual(updated_content, "Updated content\n")
def test_update_files_no_filename_single_file_in_chat(self):
sample_file = "accumulate.py"
content = (
"def accumulate(collection, operation):\n return [operation(x) for x in"
" collection]\n"
)
with open(sample_file, "w") as f:
f.write("Original content\n")
# Initialize WholeFileCoder with the temporary directory
io = InputOutput(yes=True)
coder = WholeFileCoder(main_model=self.GPT35, io=io, fnames=[sample_file])
# Set the partial response content with the updated content
coder.partial_response_content = (
f"Here's the modified `{sample_file}` file that implements the `accumulate`"
f" function as per the given instructions:\n\n```\n{content}```\n\nThis"
" implementation uses a list comprehension to apply the `operation` function to"
" each element of the `collection` and returns the resulting list."
)
# Call update_files method
edited_files = coder.apply_updates()
# Check if the sample file was updated
self.assertIn(sample_file, edited_files)
# Check if the content of the sample file was updated
with open(sample_file, "r") as f:
updated_content = f.read()
self.assertEqual(updated_content, content)
def test_update_files_earlier_filename(self):
fname_a = Path("a.txt")
fname_b = Path("b.txt")
fname_a.write_text("before a\n")
fname_b.write_text("before b\n")
response = """
Here is a new version of `a.txt` for you to consider:
```
after a
```
And here is `b.txt`:
```
after b
```
"""
# Initialize WholeFileCoder with the temporary directory
io = InputOutput(yes=True)
coder = WholeFileCoder(main_model=self.GPT35, io=io, fnames=[fname_a, fname_b])
# Set the partial response content with the updated content
coder.partial_response_content = response
# Call update_files method
edited_files = coder.apply_updates()
# Check if the sample file was updated
self.assertIn(str(fname_a), edited_files)
self.assertIn(str(fname_b), edited_files)
self.assertEqual(fname_a.read_text(), "after a\n")
self.assertEqual(fname_b.read_text(), "after b\n")
def test_update_hash_filename(self):
fname_a = Path("a.txt")
fname_b = Path("b.txt")
fname_a.write_text("before a\n")
fname_b.write_text("before b\n")
response = """
### a.txt
```
after a
```
### b.txt
```
after b
```
"""
# Initialize WholeFileCoder with the temporary directory
io = InputOutput(yes=True)
coder = WholeFileCoder(main_model=self.GPT35, io=io, fnames=[fname_a, fname_b])
# Set the partial response content with the updated content
coder.partial_response_content = response
# Call update_files method
edited_files = coder.apply_updates()
dump(edited_files)
# Check if the sample file was updated
self.assertIn(str(fname_a), edited_files)
self.assertIn(str(fname_b), edited_files)
self.assertEqual(fname_a.read_text(), "after a\n")
self.assertEqual(fname_b.read_text(), "after b\n")
def test_update_named_file_but_extra_unnamed_code_block(self):
sample_file = "hello.py"
new_content = "new\ncontent\ngoes\nhere\n"
with open(sample_file, "w") as f:
f.write("Original content\n")
# Initialize WholeFileCoder with the temporary directory
io = InputOutput(yes=True)
coder = WholeFileCoder(main_model=self.GPT35, io=io, fnames=[sample_file])
# Set the partial response content with the updated content
coder.partial_response_content = (
f"Here's the modified `{sample_file}` file that implements the `accumulate`"
f" function as per the given instructions:\n\n```\n{new_content}```\n\nThis"
" implementation uses a list comprehension to apply the `operation` function to"
" each element of the `collection` and returns the resulting list.\n"
"Run it like this:\n\n"
"```\npython {sample_file}\n```\n\n"
)
# Call update_files method
edited_files = coder.apply_updates()
# Check if the sample file was updated
self.assertIn(sample_file, edited_files)
# Check if the content of the sample file was updated
with open(sample_file, "r") as f:
updated_content = f.read()
self.assertEqual(updated_content, new_content)
def test_full_edit(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(self.GPT35, "whole", io=InputOutput(), fnames=files, stream=False)
# no trailing newline so the response content below doesn't add ANOTHER newline
new_content = "new\ntwo\nthree"
def mock_send(*args, **kwargs):
coder.partial_response_content = f"""
Do this:
{Path(file1).name}
```
{new_content}
```
"""
coder.partial_response_function_call = dict()
return []
coder.send = MagicMock(side_effect=mock_send)
# Call the run method with a message
coder.run(with_message="hi")
content = Path(file1).read_text(encoding="utf-8")
# check for one trailing newline
self.assertEqual(content, new_content + "\n")
if __name__ == "__main__":
unittest.main()