feat: Add a new source file and update the RepoMap test

This commit is contained in:
Paul Gauthier 2024-08-19 15:50:06 -07:00 committed by Paul Gauthier (aider)
parent 78a4ed69d8
commit d875e7ebe0

View file

@ -49,11 +49,14 @@ class TestRepoMap(unittest.TestCase):
# Create two source files with one function each
file1_content = "def function1():\n return 'Hello from file1'\n"
file2_content = "def function2():\n return 'Hello from file2'\n"
file3_content = "def function3():\n return 'Hello from file3'\n"
with open(os.path.join(temp_dir, "file1.py"), "w") as f:
f.write(file1_content)
with open(os.path.join(temp_dir, "file2.py"), "w") as f:
f.write(file2_content)
with open(os.path.join(temp_dir, "file3.py"), "w") as f:
f.write(file3_content)
# Add files to git
repo.index.add(["file1.py", "file2.py"])
@ -62,30 +65,30 @@ class TestRepoMap(unittest.TestCase):
# Initialize RepoMap with refresh="files" and one source file
io = InputOutput()
repo_map = RepoMap(main_model=self.GPT35, root=temp_dir, io=io, refresh="files")
chat_files = [os.path.join(temp_dir, "file1.py")]
other_files = [os.path.join(temp_dir, "file2.py")]
chat_files = [os.path.join(temp_dir, "file2.py")]
other_files = [os.path.join(temp_dir, "file1.py"), os.path.join(temp_dir, "file3.py")]
# Get initial repo map
initial_map = repo_map.get_repo_map(chat_files, other_files)
self.assertIn("function1", initial_map)
dump(initial_map)
self.assertNotIn("function2", initial_map)
# Add a 2nd function to file1.py
with open(os.path.join(temp_dir, "file1.py"), "a") as f:
f.write("\ndef function3():\n return 'Hello from function3'\n")
with open(os.path.join(temp_dir, "file1.py"), "w") as f:
f.write("\ndef functionNEW():\n return 'Hello NEW'\n")
# Get another repo map
second_map = repo_map.get_repo_map(chat_files, other_files)
self.assertEqual(initial_map, second_map, "RepoMap should not change without refresh")
# Add the 2nd file to the chat
chat_files.append(os.path.join(temp_dir, "file2.py"))
chat_files = [os.path.join(temp_dir, "file2.py"), os.path.join(temp_dir, "file3.py")]
other_files = [os.path.join(temp_dir, "file1.py")]
# Get a new repo map
final_map = repo_map.get_repo_map(chat_files, [])
self.assertIn("function1", final_map)
self.assertIn("function2", final_map)
self.assertIn("function3", final_map)
final_map = repo_map.get_repo_map(chat_files, other_files)
dump(final_map)
self.assertIn("functionNEW", final_map)
# close the open cache files, so Windows won't error
del repo_map