Fixes some bugs in test code

.github/workflows/ci.yml

@@ -43,6 +43,9 @@ jobs:
 
       - name: Install Poetry
         uses: snok/install-poetry@v1
+        with:
+          version: 1.8.1
+          plugins: poetry-plugin-export
 
       - name: build-java
         run: ./gradlew ${{ env.GRADLE_SWITCHES }} build

rewrite/pyproject.toml

@@ -20,4 +20,4 @@ openrewrite-remote = "*"
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]
-build-backend = "poetry.core.masonry.api"
+build-backend = "poetry.core.masonry.api"

rewrite/rewrite/java/extensions.py

@@ -1,6 +1,8 @@
 from typing import Optional, TypeVar, TYPE_CHECKING
 
 from .support_types import J, JRightPadded, JLeftPadded, JContainer, Space
+from ..visitor import Cursor
+from ..tree import Tree
 
 if TYPE_CHECKING:
     from .visitor import JavaVisitor
@@ -16,7 +18,7 @@ def visit_container(v: 'JavaVisitor', container: Optional[JContainer[J2]], loc:
 
     v.cursor = Cursor(v.cursor, container)
     before = v.visit_space(container.before, loc.before_location, p)
-    js = [v.visit_right_padded(el.element, loc.element_location, p) for el in container.padding.elements]
+    js = [v.visit_right_padded(el, loc.element_location, p) for el in container.padding.elements]
     v.cursor = v.cursor.parent
 
     return container if js == container.padding.elements and before is container.before else JContainer(before, js, container.markers)

rewrite/rewrite/python/support_types.py

@@ -23,6 +23,10 @@ def accept(self, v: TreeVisitor[Any, P], p: P) -> Optional[Any]:
     def accept_python(self, v: 'PythonVisitor[P]', p: P) -> Optional['J']:
         ...
 
+    def is_acceptable(self, v: TreeVisitor[Any, P], p: P) -> bool:
+        from .visitor import PythonVisitor
+        return isinstance(v, PythonVisitor)
+
 
 class PySpace:
     class Location(Enum):

rewrite/rewrite/test.py

@@ -72,7 +72,9 @@ def rewrite_run(*sources: list[SourceSpec]):
 
                     if spec.after is not None:
                         after = spec.after(source_file.print_all())
-                        assert after == spec.before
+                        assert source_file.print_all() == after
+                    else:
+                        assert source_file.print_all() == spec.before
                     break
     except Exception:
         raise
@@ -85,6 +87,6 @@ def python(before: str, after: str = None) -> list[SourceSpec]:
         random_id(),
         PythonParserBuilder(),
         textwrap.dedent(before),
-        None if after is None else lambda: after,
+        None if after is None else lambda _: textwrap.dedent(after),
         None
     )]

rewrite/rewrite/visitor.py

@@ -41,7 +41,7 @@ def first_enclosing(self, type: Type[P]) -> P:
 class TreeVisitor(Protocol[T, P]):
     _visit_count: int = 0
     _cursor: Cursor = Cursor(None, "root")
-    _after_visit: Optional[List[TreeVisitor[Any, P]]]
+    _after_visit: Optional[List[TreeVisitor[Any, P]]] = None
 
     def is_acceptable(self, source_file: SourceFile, p: P) -> bool:
         return True