Fix directory tests (#458)

Author: bagel897

.github/actions/run-ats/action.yml

@@ -48,11 +48,10 @@ runs:
           echo "No tests to run, skipping..."
           exit 0
         fi
-        echo $TESTS_TO_RUN | xargs uv run pytest --cov \
+        echo $TESTS_TO_RUN | xargs uv run --frozen pytest --cov \
           -o junit_suite_name="${{ github.job }}" \
           -n auto \
           -vv \
-          --cov \
           --cov-append \
           ${{ inputs.collect_args }}
 

src/codegen/extensions/lsp/io.py

@@ -16,7 +16,7 @@
 
 @dataclass
 class File:
-    doc: TextDocument
+    doc: TextDocument | None
     path: Path
     change: TextEdit | None = None
     other_change: CreateFile | RenameFile | DeleteFile | None = None
@@ -65,6 +65,8 @@ def read_text(self, path: Path) -> str:
             return file.change.new_text
         if file.created:
             return ""
+        if file.doc is None:
+            return self.base_io.read_text(path)
         return file.doc.source
 
     def read_bytes(self, path: Path) -> bytes:
@@ -76,6 +78,8 @@ def read_bytes(self, path: Path) -> bytes:
             return file.change.new_text.encode("utf-8")
         if file.created:
             return b""
+        if file.doc is None:
+            return self.base_io.read_bytes(path)
         return file.doc.source.encode("utf-8")
 
     def write_bytes(self, path: Path, content: bytes) -> None:
@@ -112,6 +116,8 @@ def file_exists(self, path: Path) -> bool:
             return True
         if file.created:
             return True
+        if file.doc is None:
+            return self.base_io.file_exists(path)
         try:
             file.doc.source
             return True
@@ -134,3 +140,13 @@ def get_workspace_edit(self) -> types.WorkspaceEdit:
                 file.change = None
         logger.info(f"Workspace edit: {pprint.pformat(list(map(asdict, document_changes)))}")
         return types.WorkspaceEdit(document_changes=document_changes)
+
+    def update_file(self, path: Path, version: int | None = None) -> None:
+        file = self._get_file(path)
+        file.doc = self.workspace.get_text_document(path.as_uri())
+        if version is not None:
+            file.version = version
+
+    def close_file(self, path: Path) -> None:
+        file = self._get_file(path)
+        file.doc = None

src/codegen/extensions/lsp/lsp.py

@@ -24,6 +24,7 @@ def did_open(server: CodegenLanguageServer, params: types.DidOpenTextDocumentPar
     # The document is automatically added to the workspace by pygls
     # We can perform any additional processing here if needed
     path = get_path(params.text_document.uri)
+    server.io.update_file(path, params.text_document.version)
     file = server.codebase.get_file(str(path), optional=True)
     if not isinstance(file, SourceFile) and path.suffix in server.codebase.ctx.extensions:
         sync = DiffLite(change_type=ChangeType.Added, path=path)
@@ -37,6 +38,7 @@ def did_change(server: CodegenLanguageServer, params: types.DidChangeTextDocumen
     # The document is automatically updated in the workspace by pygls
     # We can perform any additional processing here if needed
     path = get_path(params.text_document.uri)
+    server.io.update_file(path, params.text_document.version)
     sync = DiffLite(change_type=ChangeType.Modified, path=path)
     server.codebase.ctx.apply_diffs([sync])
 
@@ -63,6 +65,8 @@ def did_close(server: CodegenLanguageServer, params: types.DidCloseTextDocumentP
     logger.info(f"Document closed: {params.text_document.uri}")
     # The document is automatically removed from the workspace by pygls
     # We can perform any additional cleanup here if needed
+    path = get_path(params.text_document.uri)
+    server.io.close_file(path)
 
 
 @server.feature(

src/codegen/sdk/core/codebase.py

@@ -510,7 +510,7 @@ def get_file_from_path(path: Path) -> File | None:
         if file is not None:
             return file
         absolute_path = self.ctx.to_absolute(filepath)
-        if absolute_path.suffix in self.ctx.extensions:
+        if absolute_path.suffix in self.ctx.extensions and not self.ctx.io.file_exists(absolute_path):
             return None
         if self.ctx.io.file_exists(absolute_path):
             return get_file_from_path(absolute_path)

tests/unit/codegen/extensions/lsp/test_workspace_sync.py

@@ -1,3 +1,5 @@
+from typing import Callable
+
 import pytest
 from lsprotocol.types import (
     DidChangeTextDocumentParams,
@@ -16,6 +18,7 @@
 from pytest_lsp import LanguageClient
 
 from codegen.sdk.core.codebase import Codebase
+from tests.unit.codegen.extensions.lsp.utils import apply_edit
 
 
 @pytest.fixture()
@@ -88,6 +91,7 @@ def example_function():
             """.strip(),
         ),
     ],
+    ids=["example_function"],
     indirect=["document_uri", "original"],
 )
 async def test_did_change(
@@ -137,7 +141,7 @@ def example_function():
     pass
                 """.strip(),
             },
-            "file://{worskpaceFolder}test.py",
+            "file://{worskpaceFolder}/test.py",
         ),
     ],
 )
@@ -165,11 +169,11 @@ async def test_did_close(
 
     # Verify the document is removed from the workspace
     document = await client.workspace_text_document_content_async(TextDocumentContentParams(uri=document_uri))
-    assert document.text == ""
+    assert document.text == original["test.py"]
 
 
 @pytest.mark.parametrize(
-    "original, document_uri, position, new_name, expected_text",
+    "original, document_uri, position, new_name, expected",
     [
         (
             {
@@ -182,15 +186,17 @@ def main():
                 """.strip(),
             },
             "file://{workspaceFolder}/test.py",
-            Position(line=0, character=0),  # Position of 'example_function'
+            Position(line=0, character=5),  # Position of 'example_function'
             "renamed_function",
-            """
+            {
+                "test.py": """
 def renamed_function():
     pass # modified
 
 def main():
     renamed_function()
                 """.strip(),
+            },
         ),
     ],
     indirect=["document_uri", "original"],
@@ -202,7 +208,7 @@ async def test_rename_after_sync(
     document_uri: str,
     position: Position,
     new_name: str,
-    expected_text: str,
+    assert_expected: Callable,
 ):
     # First open the document
     client.text_document_did_open(
@@ -243,8 +249,6 @@ async def test_rename_after_sync(
             new_name=new_name,
         )
     )
-
-    # Verify the rename was successful
-    document = await client.workspace_text_document_content_async(TextDocumentContentParams(uri=document_uri))
-    assert document is not None
-    assert document.text == expected_text
+    if result:
+        apply_edit(codebase, result)
+    assert_expected(codebase)

tests/unit/codegen/extensions/lsp/utils.py

@@ -10,7 +10,6 @@ def apply_edit(codebase: Codebase, edit: WorkspaceEdit):
             path = get_path(change.text_document.uri)
             file = codebase.get_file(str(path.relative_to(codebase.repo_path)))
             for edit in change.edits:
-                print("BRUH")
                 file.edit(edit.new_text)
         if isinstance(change, CreateFile):
             path = get_path(change.uri)

tests/unit/codegen/sdk/core/test_directory.py

@@ -12,7 +12,7 @@
 
 
 @pytest.fixture
-def mock_codebase_graph(tmp_path):
+def mock_codebase_context(tmp_path):
     mock = MagicMock(spec=CodebaseContext)
     mock.transaction_manager = MagicMock()
     mock.config = CodebaseConfig()
@@ -39,8 +39,8 @@ def sub_dir(subdir_path, tmp_path):
 
 
 @pytest.fixture
-def mock_file(dir_path, mock_codebase_graph):
-    return File(filepath=dir_path / "example.py", ctx=mock_codebase_graph)
+def mock_file(dir_path, mock_codebase_context):
+    return File(filepath=dir_path / "example.py", ctx=mock_codebase_context)
 
 
 @pytest.fixture
@@ -66,9 +66,9 @@ def test_name_property(mock_directory):
     assert mock_directory.name == "mock_dir"
 
 
-def test_add_and_file(mock_directory, mock_codebase_graph):
+def test_add_and_file(mock_directory, mock_codebase_context):
     """Test adding a file to the directory."""
-    mock_file = File(filepath=Path("mock_dir/example_2.py"), ctx=mock_codebase_graph)
+    mock_file = File(filepath=Path("mock_dir/example_2.py"), ctx=mock_codebase_context)
     mock_directory.add_file(mock_file)
     rel_path = os.path.relpath(mock_file.file_path, mock_directory.dirpath)
     assert rel_path in mock_directory.items
@@ -139,12 +139,12 @@ def test_get_subdirectory(mock_directory, sub_dir):
     assert retrieved_subdir is sub_dir
 
 
-def test_files_property(mock_directory, sub_dir, mock_codebase_graph):
+def test_files_property(mock_directory, sub_dir, mock_codebase_context):
     """Test the 'files' property returns all files recursively."""
     all_files = mock_directory.files
     assert len(all_files) == 1
 
-    new_file = File(filepath=Path("mock_dir/example_2.py"), ctx=mock_codebase_graph)
+    new_file = File(filepath=Path("mock_dir/example_2.py"), ctx=mock_codebase_context)
     sub_dir.add_file(new_file)
 
     all_files = mock_directory.files
@@ -171,26 +171,26 @@ def test_subdirectories_property(mock_directory, sub_dir):
     assert new_sub_dir in all_subdirs
 
 
-def test_update_filepath(mock_directory, mock_codebase_graph, mock_file):
+def test_update_filepath(mock_directory, mock_codebase_context, mock_file):
     """Test updating file paths when the directory path changes."""
     mock_directory.update_filepath("/absolute/new_mock_dir")
 
     # Verify the files have updated file paths
-    mock_codebase_graph.transaction_manager.add_file_rename_transaction.assert_called_once_with(mock_file, "/absolute/new_mock_dir/example.py")
+    mock_codebase_context.transaction_manager.add_file_rename_transaction.assert_called_once_with(mock_file, "/absolute/new_mock_dir/example.py")
 
 
-def test_remove(mock_directory, sub_dir, mock_codebase_graph, mock_file):
+def test_remove(mock_directory, sub_dir, mock_codebase_context, mock_file):
     mock_directory.remove()
 
-    mock_codebase_graph.transaction_manager.add_file_remove_transaction.assert_called_once_with(mock_file)
+    mock_codebase_context.transaction_manager.add_file_remove_transaction.assert_called_once_with(mock_file)
 
 
-def test_rename(mock_directory, mock_codebase_graph, mock_file):
+def test_rename(mock_directory, mock_codebase_context, mock_file):
     """Test renaming the directory."""
     mock_directory.rename("renamed_dir")
     # This fails because it is not implemented to rename the directory itself.
     # assert mock_directory.dirpath == "/absolute/renamed_dir"
-    mock_codebase_graph.transaction_manager.add_file_rename_transaction.assert_called_once_with(mock_file, "renamed_dir/example.py")
+    mock_codebase_context.transaction_manager.add_file_rename_transaction.assert_called_once_with(mock_file, "renamed_dir/example.py")
 
 
 def test_iteration(mock_directory):