Try to update test_improve_existing_project

tests/applications/cli/test_main.py

@@ -1,9 +1,15 @@
 # Generated by CodiumAI
+from pathlib import Path
+from typing import List
+
+import pytest
+import toml
 
 import gpt_engineer.applications.cli.main as main
 
 from gpt_engineer.core.default.disk_execution_env import DiskExecutionEnv
-from gpt_engineer.core.default.paths import ENTRYPOINT_FILE
+from gpt_engineer.core.default.disk_memory import DiskMemory
+from gpt_engineer.core.default.paths import ENTRYPOINT_FILE, metadata_path
 from tests.caching_ai import CachingAI
 
 main.AI = CachingAI
@@ -47,6 +53,24 @@ def input_generator():
 prompt_text = "Make a python program that writes 'hello' to a file called 'output.txt'"
 
 
+@pytest.fixture
+def mock_editor_file_selector(monkeypatch):
+    """Mock the editor file selector to simulate user interaction."""
+
+    def mock(input_path: str, init: bool = True) -> List[str]:
+        toml_file = DiskMemory(metadata_path(input_path)).path / "file_selection.toml"
+        tree_dict = {
+            "files": {str(Path(input_path) / "example_file.py"): {"selected": True}}
+        }
+        with open(toml_file, "w") as f:
+            toml.dump(tree_dict, f)
+        return [str(Path(input_path) / "example_file.py")]
+
+    monkeypatch.setattr(
+        "gpt_engineer.applications.cli.file_selector.editor_file_selector", mock
+    )
+
+
 class TestMain:
     #  Runs gpt-engineer with default settings and generates a project in the specified path.
     def test_default_settings_generate_project(self, tmp_path, monkeypatch):
@@ -63,40 +87,21 @@ def test_default_settings_generate_project(self, tmp_path, monkeypatch):
         assert text == "hello"
 
     #  Runs gpt-engineer with improve mode and improves an existing project in the specified path.
-    def test_improve_existing_project(self, tmp_path, monkeypatch):
-        from pathlib import Path
-        from typing import List
-
-        import toml
-
-        from gpt_engineer.core.default.disk_memory import DiskMemory
-        from gpt_engineer.core.default.paths import metadata_path
-
-        def mock_editor_file_selector(input_path: str, init: bool = True) -> List[str]:
-            toml_file = (
-                DiskMemory(metadata_path(input_path)).path / "file_selection.toml"
-            )
-            tree_dict = {"files": {"example_file.py": {"selected": True}}}
-
-            with open(toml_file, "w") as f:
-                toml.dump(tree_dict, f)
-
-            return [str(Path(input_path) / "example_file.py")]
-
-        monkeypatch.setattr(
-            "gpt_engineer.applications.cli.file_selector.editor_file_selector",
-            mock_editor_file_selector,
+    def test_improve_existing_project(self, tmp_path, mock_editor_file_selector):
+        project_dir = tmp_path / "projects/example"
+        project_dir.mkdir(parents=True)
+        (project_dir / "prompt").write_text(
+            "Make a python program that writes 'hello' to a file called 'output.txt'"
         )
 
-        p = tmp_path / "projects/example"
-        p.mkdir(parents=True)
-        (p / "prompt").write_text(prompt_text)
-        simplified_main(str(p), "improve")
+        # Run the main function in improve mode
+        simplified_main(str(project_dir), "improve")
 
-        ex_env = DiskExecutionEnv(path=p)
+        ex_env = DiskExecutionEnv(path=project_dir)
         ex_env.run(f"bash {ENTRYPOINT_FILE}")
-        assert (p / "output.txt").exists()
-        text = (p / "output.txt").read_text().strip()
+
+        assert (project_dir / "output.txt").exists()
+        text = (project_dir / "output.txt").read_text().strip()
         assert text == "hello"
 
     #  Runs gpt-engineer with lite mode and generates a project with only the main prompt.