From 3401395e61a56dedd07c0b0c1cfa92b78cf91352 Mon Sep 17 00:00:00 2001
From: Paul Gauthier <aider@paulg.org>
Date: Thu, 17 Aug 2023 11:08:50 -0700
Subject: [PATCH] fix wholefile tests

---
 aider/coders/base_coder.py | 16 ++++++++--------
 tests/test_wholefile.py    |  2 +-
 2 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/aider/coders/base_coder.py b/aider/coders/base_coder.py
index be55ada6c..e5cc3f2a0 100755
--- a/aider/coders/base_coder.py
+++ b/aider/coders/base_coder.py
@@ -801,13 +801,17 @@ class Coder:
 
         return res
 
+    def update_files(self):
+        edits = self.get_edits()
+        edits = self.prepare_to_edit(edits)
+        self.apply_edits(edits)
+        return set(edit[0] for edit in edits)
+
     def apply_updates(self):
         max_apply_update_errors = 3
 
         try:
-            edits = self.get_edits()
-            edits = self.prepare_to_edit(edits)
-            self.apply_edits(edits)
+            edited = self.update_files()
         except ValueError as err:
             err = err.args[0]
             self.apply_update_errors += 1
@@ -833,11 +837,7 @@ class Coder:
 
         self.apply_update_errors = 0
 
-        # TODO FIXME: make sure
-        edited = set()
-        for edit in sorted(edits):
-            path = edit[0]
-            edited.add(path)
+        for path in edited:
             if self.dry_run:
                 self.io.tool_output(f"Did not apply edit to {path} (--dry-run)")
             else:
diff --git a/tests/test_wholefile.py b/tests/test_wholefile.py
index 8f9f89fc5..29cc74fa2 100644
--- a/tests/test_wholefile.py
+++ b/tests/test_wholefile.py
@@ -90,7 +90,7 @@ class TestWholeFileCoder(unittest.TestCase):
         # Set the partial response content with the updated content
         coder.partial_response_content = f"{sample_file}\n```\n0\n\1\n2\n"
 
-        lines = coder.update_files(mode="diff").splitlines()
+        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)