Snapshot tests for editing, WIP

Need to look over the snapshot results to see whether they're what we expect. Also, for some reason tests are failing...

pyproject.toml

@@ -56,6 +56,7 @@ dev = [
   "ruff>=0.0.285",
   "toml-sort>=0.24.2",
   "types-pyyaml>=6.0.12.20241230",
+  "syrupy>=4.9.1",
 ]
 agno = [
     "agno>=1.2.6",

test/serena/__snapshots__/test_symbol_editing.ambr

@@ -0,0 +1,700 @@
+# serializer version: 1
+# name: test_delete_symbol[test_case0]
+  '''
+  """
+  Test module for variable declarations and usage.
+  
+  This module tests various types of variable declarations and usages including:
+  - Module-level variables
+  - Class-level variables
+  - Instance variables
+  - Variable reassignments
+  """
+  
+  from dataclasses import dataclass, field
+  
+  # Module-level variables
+  module_var = "Initial module value"
+  
+  reassignable_module_var = 10
+  reassignable_module_var = 20  # Reassigned
+  
+  # Module-level variable with type annotation
+  typed_module_var: int = 42
+  
+  
+  # Regular class with class and instance variables
+  
+  
+  
+  # Dataclass with variables
+  @dataclass
+  class VariableDataclass:
+      """Dataclass that contains various fields."""
+  
+      # Field variables with type annotations
+      id: int
+      name: str
+      items: list[str] = field(default_factory=list)
+      metadata: dict[str, str] = field(default_factory=dict)
+      optional_value: float | None = None
+  
+      # This will be reassigned in various places
+      status: str = "pending"
+  
+  
+  # Function that uses the module variables
+  def use_module_variables():
+      """Function that uses module-level variables."""
+      result = module_var + " used in function"
+      other_result = reassignable_module_var * 2
+      return result, other_result
+  
+  
+  # Create instances and use variables
+  dataclass_instance = VariableDataclass(id=1, name="Test")
+  dataclass_instance.status = "active"  # Reassign dataclass field
+  
+  # Use variables at module level
+  module_result = module_var + " used at module level"
+  other_module_result = reassignable_module_var + 30
+  
+  # Create a second dataclass instance with different status
+  second_dataclass = VariableDataclass(id=2, name="Another Test")
+  second_dataclass.status = "completed"  # Another reassignment of status
+  
+  '''
+# ---
+# name: test_delete_symbol[test_case1]
+  '''
+  
+  
+  export function helperFunction() {
+      const demo = new DemoClass(42);
+      demo.printValue();
+  }
+  
+  helperFunction();
+  
+  '''
+# ---
+# name: test_insert_in_rel_to_symbol[test_case0-after]
+  '''
+  """
+  Test module for variable declarations and usage.
+  
+  This module tests various types of variable declarations and usages including:
+  - Module-level variables
+  - Class-level variables
+  - Instance variables
+  - Variable reassignments
+  """
+  
+  from dataclasses import dataclass, field
+  
+  # Module-level variables
+  module_var = "Initial module value"
+  
+  reassignable_module_var = 10
+  reassignable_module_var = 20  # Reassigned
+  
+  # Module-level variable with type annotation
+  typed_module_var: int = 42
+  new_module_var = "Inserted after typed_module_var"
+  
+  # Regular class with class and instance variables
+  class VariableContainer:
+      """Class that contains various variables."""
+  
+      # Class-level variables
+      class_var = "Initial class value"
+  
+      reassignable_class_var = True
+      reassignable_class_var = False  # Reassigned #noqa: PIE794
+  
+      # Class-level variable with type annotation
+      typed_class_var: str = "typed value"
+  
+      def __init__(self):
+          # Instance variables
+          self.instance_var = "Initial instance value"
+          self.reassignable_instance_var = 100
+  
+          # Instance variable with type annotation
+          self.typed_instance_var: list[str] = ["item1", "item2"]
+  
+      def modify_instance_var(self):
+          # Reassign instance variable
+          self.instance_var = "Modified instance value"
+          self.reassignable_instance_var = 200  # Reassigned
+  
+      def use_module_var(self):
+          # Use module-level variables
+          result = module_var + " used in method"
+          other_result = reassignable_module_var + 5
+          return result, other_result
+  
+      def use_class_var(self):
+          # Use class-level variables
+          result = VariableContainer.class_var + " used in method"
+          other_result = VariableContainer.reassignable_class_var
+          return result, other_result
+  
+  
+  # Dataclass with variables
+  @dataclass
+  class VariableDataclass:
+      """Dataclass that contains various fields."""
+  
+      # Field variables with type annotations
+      id: int
+      name: str
+      items: list[str] = field(default_factory=list)
+      metadata: dict[str, str] = field(default_factory=dict)
+      optional_value: float | None = None
+  
+      # This will be reassigned in various places
+      status: str = "pending"
+  
+  
+  # Function that uses the module variables
+  def use_module_variables():
+      """Function that uses module-level variables."""
+      result = module_var + " used in function"
+      other_result = reassignable_module_var * 2
+      return result, other_result
+  
+  
+  # Create instances and use variables
+  dataclass_instance = VariableDataclass(id=1, name="Test")
+  dataclass_instance.status = "active"  # Reassign dataclass field
+  
+  # Use variables at module level
+  module_result = module_var + " used at module level"
+  other_module_result = reassignable_module_var + 30
+  
+  # Create a second dataclass instance with different status
+  second_dataclass = VariableDataclass(id=2, name="Another Test")
+  second_dataclass.status = "completed"  # Another reassignment of status
+  
+  '''
+# ---
+# name: test_insert_in_rel_to_symbol[test_case0-before]
+  '''
+  """
+  Test module for variable declarations and usage.
+  
+  This module tests various types of variable declarations and usages including:
+  - Module-level variables
+  - Class-level variables
+  - Instance variables
+  - Variable reassignments
+  """
+  
+  from dataclasses import dataclass, field
+  
+  # Module-level variables
+  module_var = "Initial module value"
+  
+  reassignable_module_var = 10
+  reassignable_module_var = 20  # Reassigned
+  
+  new_module_var = "Inserted after typed_module_var"# Module-level variable with type annotation
+  typed_module_var: int = 42
+  
+  
+  # Regular class with class and instance variables
+  class VariableContainer:
+      """Class that contains various variables."""
+  
+      # Class-level variables
+      class_var = "Initial class value"
+  
+      reassignable_class_var = True
+      reassignable_class_var = False  # Reassigned #noqa: PIE794
+  
+      # Class-level variable with type annotation
+      typed_class_var: str = "typed value"
+  
+      def __init__(self):
+          # Instance variables
+          self.instance_var = "Initial instance value"
+          self.reassignable_instance_var = 100
+  
+          # Instance variable with type annotation
+          self.typed_instance_var: list[str] = ["item1", "item2"]
+  
+      def modify_instance_var(self):
+          # Reassign instance variable
+          self.instance_var = "Modified instance value"
+          self.reassignable_instance_var = 200  # Reassigned
+  
+      def use_module_var(self):
+          # Use module-level variables
+          result = module_var + " used in method"
+          other_result = reassignable_module_var + 5
+          return result, other_result
+  
+      def use_class_var(self):
+          # Use class-level variables
+          result = VariableContainer.class_var + " used in method"
+          other_result = VariableContainer.reassignable_class_var
+          return result, other_result
+  
+  
+  # Dataclass with variables
+  @dataclass
+  class VariableDataclass:
+      """Dataclass that contains various fields."""
+  
+      # Field variables with type annotations
+      id: int
+      name: str
+      items: list[str] = field(default_factory=list)
+      metadata: dict[str, str] = field(default_factory=dict)
+      optional_value: float | None = None
+  
+      # This will be reassigned in various places
+      status: str = "pending"
+  
+  
+  # Function that uses the module variables
+  def use_module_variables():
+      """Function that uses module-level variables."""
+      result = module_var + " used in function"
+      other_result = reassignable_module_var * 2
+      return result, other_result
+  
+  
+  # Create instances and use variables
+  dataclass_instance = VariableDataclass(id=1, name="Test")
+  dataclass_instance.status = "active"  # Reassign dataclass field
+  
+  # Use variables at module level
+  module_result = module_var + " used at module level"
+  other_module_result = reassignable_module_var + 30
+  
+  # Create a second dataclass instance with different status
+  second_dataclass = VariableDataclass(id=2, name="Another Test")
+  second_dataclass.status = "completed"  # Another reassignment of status
+  
+  '''
+# ---
+# name: test_insert_in_rel_to_symbol[test_case1-after]
+  '''
+  """
+  Test module for variable declarations and usage.
+  
+  This module tests various types of variable declarations and usages including:
+  - Module-level variables
+  - Class-level variables
+  - Instance variables
+  - Variable reassignments
+  """
+  
+  from dataclasses import dataclass, field
+  
+  # Module-level variables
+  module_var = "Initial module value"
+  
+  reassignable_module_var = 10
+  reassignable_module_var = 20  # Reassigned
+  
+  # Module-level variable with type annotation
+  typed_module_var: int = 42
+  
+  
+  # Regular class with class and instance variables
+  class VariableContainer:
+      """Class that contains various variables."""
+  
+      # Class-level variables
+      class_var = "Initial class value"
+  
+      reassignable_class_var = True
+      reassignable_class_var = False  # Reassigned #noqa: PIE794
+  
+      # Class-level variable with type annotation
+      typed_class_var: str = "typed value"
+  
+      def __init__(self):
+          # Instance variables
+          self.instance_var = "Initial instance value"
+          self.reassignable_instance_var = 100
+  
+          # Instance variable with type annotation
+          self.typed_instance_var: list[str] = ["item1", "item2"]
+  
+      def modify_instance_var(self):
+          # Reassign instance variable
+          self.instance_var = "Modified instance value"
+          self.reassignable_instance_var = 200  # Reassigned
+  
+      def use_module_var(self):
+          # Use module-level variables
+          result = module_var + " used in method"
+          other_result = reassignable_module_var + 5
+          return result, other_result
+  
+      def use_class_var(self):
+          # Use class-level variables
+          result = VariableContainer.class_var + " used in method"
+          other_result = VariableContainer.reassignable_class_var
+          return result, other_result
+  
+  
+  # Dataclass with variables
+  @dataclass
+  class VariableDataclass:
+      """Dataclass that contains various fields."""
+  
+      # Field variables with type annotations
+      id: int
+      name: str
+      items: list[str] = field(default_factory=list)
+      metadata: dict[str, str] = field(default_factory=dict)
+      optional_value: float | None = None
+  
+      # This will be reassigned in various places
+      status: str = "pending"
+  
+  
+  # Function that uses the module variables
+  def use_module_variables():
+      """Function that uses module-level variables."""
+      result = module_var + " used in function"
+      other_result = reassignable_module_var * 2
+      return result, other_result
+  def new_inserted_function():
+      print("This is a new function inserted before another.")
+  
+  # Create instances and use variables
+  dataclass_instance = VariableDataclass(id=1, name="Test")
+  dataclass_instance.status = "active"  # Reassign dataclass field
+  
+  # Use variables at module level
+  module_result = module_var + " used at module level"
+  other_module_result = reassignable_module_var + 30
+  
+  # Create a second dataclass instance with different status
+  second_dataclass = VariableDataclass(id=2, name="Another Test")
+  second_dataclass.status = "completed"  # Another reassignment of status
+  
+  '''
+# ---
+# name: test_insert_in_rel_to_symbol[test_case1-before]
+  '''
+  """
+  Test module for variable declarations and usage.
+  
+  This module tests various types of variable declarations and usages including:
+  - Module-level variables
+  - Class-level variables
+  - Instance variables
+  - Variable reassignments
+  """
+  
+  from dataclasses import dataclass, field
+  
+  # Module-level variables
+  module_var = "Initial module value"
+  
+  reassignable_module_var = 10
+  reassignable_module_var = 20  # Reassigned
+  
+  # Module-level variable with type annotation
+  typed_module_var: int = 42
+  
+  
+  # Regular class with class and instance variables
+  class VariableContainer:
+      """Class that contains various variables."""
+  
+      # Class-level variables
+      class_var = "Initial class value"
+  
+      reassignable_class_var = True
+      reassignable_class_var = False  # Reassigned #noqa: PIE794
+  
+      # Class-level variable with type annotation
+      typed_class_var: str = "typed value"
+  
+      def __init__(self):
+          # Instance variables
+          self.instance_var = "Initial instance value"
+          self.reassignable_instance_var = 100
+  
+          # Instance variable with type annotation
+          self.typed_instance_var: list[str] = ["item1", "item2"]
+  
+      def modify_instance_var(self):
+          # Reassign instance variable
+          self.instance_var = "Modified instance value"
+          self.reassignable_instance_var = 200  # Reassigned
+  
+      def use_module_var(self):
+          # Use module-level variables
+          result = module_var + " used in method"
+          other_result = reassignable_module_var + 5
+          return result, other_result
+  
+      def use_class_var(self):
+          # Use class-level variables
+          result = VariableContainer.class_var + " used in method"
+          other_result = VariableContainer.reassignable_class_var
+          return result, other_result
+  
+  
+  # Dataclass with variables
+  @dataclass
+  class VariableDataclass:
+      """Dataclass that contains various fields."""
+  
+      # Field variables with type annotations
+      id: int
+      name: str
+      items: list[str] = field(default_factory=list)
+      metadata: dict[str, str] = field(default_factory=dict)
+      optional_value: float | None = None
+  
+      # This will be reassigned in various places
+      status: str = "pending"
+  
+  
+  def new_inserted_function():
+      print("This is a new function inserted before another.")# Function that uses the module variables
+  def use_module_variables():
+      """Function that uses module-level variables."""
+      result = module_var + " used in function"
+      other_result = reassignable_module_var * 2
+      return result, other_result
+  
+  
+  # Create instances and use variables
+  dataclass_instance = VariableDataclass(id=1, name="Test")
+  dataclass_instance.status = "active"  # Reassign dataclass field
+  
+  # Use variables at module level
+  module_result = module_var + " used at module level"
+  other_module_result = reassignable_module_var + 30
+  
+  # Create a second dataclass instance with different status
+  second_dataclass = VariableDataclass(id=2, name="Another Test")
+  second_dataclass.status = "completed"  # Another reassignment of status
+  
+  '''
+# ---
+# name: test_insert_in_rel_to_symbol[test_case2-after]
+  '''
+  export class DemoClass {
+      value: number;
+      constructor(value: number) {
+          this.value = value;
+      }
+      printValue() {
+          console.log(this.value);
+      }
+  }
+  function newFunctionAfterClass(): void {
+      console.log("This function is after DemoClass.");
+  # }
+  export function helperFunction() {
+      const demo = new DemoClass(42);
+      demo.printValue();
+  }
+  
+  helperFunction();
+  
+  '''
+# ---
+# name: test_insert_in_rel_to_symbol[test_case2-before]
+  '''
+  function newFunctionAfterClass(): void {
+      console.log("This function is after DemoClass.");
+  # }export class DemoClass {
+      value: number;
+      constructor(value: number) {
+          this.value = value;
+      }
+      printValue() {
+          console.log(this.value);
+      }
+  }
+  
+  export function helperFunction() {
+      const demo = new DemoClass(42);
+      demo.printValue();
+  }
+  
+  helperFunction();
+  
+  '''
+# ---
+# name: test_insert_in_rel_to_symbol[test_case3-after]
+  '''
+  export class DemoClass {
+      value: number;
+      constructor(value: number) {
+          this.value = value;
+      }
+      printValue() {
+          console.log(this.value);
+      }
+  }
+  
+  export function helperFunction() {
+      const demo = new DemoClass(42);
+      demo.printValue();
+  }
+  function newInsertedFunction(): void {
+      console.log("This is a new function inserted before another.");
+  }
+  helperFunction();
+  
+  '''
+# ---
+# name: test_insert_in_rel_to_symbol[test_case3-before]
+  '''
+  export class DemoClass {
+      value: number;
+      constructor(value: number) {
+          this.value = value;
+      }
+      printValue() {
+          console.log(this.value);
+      }
+  }
+  function newInsertedFunction(): void {
+      console.log("This is a new function inserted before another.");
+  }
+  export function helperFunction() {
+      const demo = new DemoClass(42);
+      demo.printValue();
+  }
+  
+  helperFunction();
+  
+  '''
+# ---
+# name: test_replace_body[test_case0]
+  '''
+  """
+  Test module for variable declarations and usage.
+  
+  This module tests various types of variable declarations and usages including:
+  - Module-level variables
+  - Class-level variables
+  - Instance variables
+  - Variable reassignments
+  """
+  
+  from dataclasses import dataclass, field
+  
+  # Module-level variables
+  module_var = "Initial module value"
+  
+  reassignable_module_var = 10
+  reassignable_module_var = 20  # Reassigned
+  
+  # Module-level variable with type annotation
+  typed_module_var: int = 42
+  
+  
+  # Regular class with class and instance variables
+  class VariableContainer:
+      """Class that contains various variables."""
+  
+      # Class-level variables
+      class_var = "Initial class value"
+  
+      reassignable_class_var = True
+      reassignable_class_var = False  # Reassigned #noqa: PIE794
+  
+      # Class-level variable with type annotation
+      typed_class_var: str = "typed value"
+  
+      def __init__(self):
+          # Instance variables
+          self.instance_var = "Initial instance value"
+          self.reassignable_instance_var = 100
+  
+          # Instance variable with type annotation
+          self.typed_instance_var: list[str] = ["item1", "item2"]
+  
+              # This body has been replaced
+          self.instance_var = "Replaced!"
+          self.reassignable_instance_var = 999
+    # Reassigned
+  
+      def use_module_var(self):
+          # Use module-level variables
+          result = module_var + " used in method"
+          other_result = reassignable_module_var + 5
+          return result, other_result
+  
+      def use_class_var(self):
+          # Use class-level variables
+          result = VariableContainer.class_var + " used in method"
+          other_result = VariableContainer.reassignable_class_var
+          return result, other_result
+  
+  
+  # Dataclass with variables
+  @dataclass
+  class VariableDataclass:
+      """Dataclass that contains various fields."""
+  
+      # Field variables with type annotations
+      id: int
+      name: str
+      items: list[str] = field(default_factory=list)
+      metadata: dict[str, str] = field(default_factory=dict)
+      optional_value: float | None = None
+  
+      # This will be reassigned in various places
+      status: str = "pending"
+  
+  
+  # Function that uses the module variables
+  def use_module_variables():
+      """Function that uses module-level variables."""
+      result = module_var + " used in function"
+      other_result = reassignable_module_var * 2
+      return result, other_result
+  
+  
+  # Create instances and use variables
+  dataclass_instance = VariableDataclass(id=1, name="Test")
+  dataclass_instance.status = "active"  # Reassign dataclass field
+  
+  # Use variables at module level
+  module_result = module_var + " used at module level"
+  other_module_result = reassignable_module_var + 30
+  
+  # Create a second dataclass instance with different status
+  second_dataclass = VariableDataclass(id=2, name="Another Test")
+  second_dataclass.status = "completed"  # Another reassignment of status
+  
+  '''
+# ---
+# name: test_replace_body[test_case1]
+  '''
+  export class DemoClass {
+      value: number;
+      constructor(value: number) {
+          this.value = value;
+      }
+              // This body has been replaced
+          console.warn("New value: " + this.value);
+  
+  }
+  
+  export function helperFunction() {
+      const demo = new DemoClass(42);
+      demo.printValue();
+  }
+  
+  helperFunction();
+  
+  '''
+# ---

test/serena/test_symbol_editing.py

@@ -5,6 +5,7 @@ from abc import abstractmethod
 from collections.abc import Iterator
 from contextlib import contextmanager
 from pathlib import Path
+from typing import Literal
 
 import pytest
 
@@ -30,35 +31,38 @@ class EditingTest:
         """Context manager for setup/teardown with a temporary directory, providing the symbol manager."""
         temp_dir = Path(tempfile.mkdtemp())
         self.repo_path = temp_dir / self.original_repo_path.name
+        language_server = None  # Initialize language_server
         try:
             shutil.copytree(self.original_repo_path, self.repo_path)
             language_server = create_ls(self.language, str(self.repo_path))
             language_server.start()
             yield SymbolManager(lang_server=language_server)
         finally:
+            if language_server is not None and language_server.is_running():
+                language_server.stop()
             shutil.rmtree(temp_dir)
 
     def _read_file(self, rel_path: str) -> str:
         """Read the content of a file in the test repository."""
+        assert self.repo_path is not None
         file_path = self.repo_path / rel_path
         with open(file_path, encoding="utf-8") as f:
             return f.read()
 
-    def run_test(self) -> None:
+    def run_test(self, content_after_ground_truth: str) -> None:
         with self._setup() as symbol_manager:
             content_before = self._read_file(self.rel_path)
             self._apply_edit(symbol_manager)
             content_after = self._read_file(self.rel_path)
             code_diff = CodeDiff(self.rel_path, original_content=content_before, modified_content=content_after)
-            self._test_diff(code_diff)
+            self._test_diff(code_diff, content_after_ground_truth)
 
     @abstractmethod
     def _apply_edit(self, symbol_manager: SymbolManager) -> None:
         pass
 
-    @abstractmethod
-    def _test_diff(self, code_diff: CodeDiff) -> None:
-        pass
+    def _test_diff(self, code_diff: CodeDiff, snapshot) -> None:
+        assert code_diff.modified_content == snapshot
 
 
 # Python test file path
@@ -67,73 +71,16 @@ PYTHON_TEST_REL_FILE_PATH = os.path.join("test_repo", "variables.py")
 # TypeScript test file path
 TYPESCRIPT_TEST_FILE = "index.ts"
 
-# Expected deleted lines for Python VariableContainer
-EXPECTED_DELETED_VARIABLE_CONTAINER_PYTHON = '''class VariableContainer:
-    """Class that contains various variables."""
-
-    # Class-level variables
-    class_var = "Initial class value"
-
-    reassignable_class_var = True
-    reassignable_class_var = False  # Reassigned #noqa: PIE794
-
-    # Class-level variable with type annotation
-    typed_class_var: str = "typed value"
-
-    def __init__(self):
-        # Instance variables
-        self.instance_var = "Initial instance value"
-        self.reassignable_instance_var = 100
-
-        # Instance variable with type annotation
-        self.typed_instance_var: list[str] = ["item1", "item2"]
-
-    def modify_instance_var(self):
-        # Reassign instance variable
-        self.instance_var = "Modified instance value"
-        self.reassignable_instance_var = 200  # Reassigned
-
-    def use_module_var(self):
-        # Use module-level variables
-        result = module_var + " used in method"
-        other_result = reassignable_module_var + 5
-        return result, other_result
-
-    def use_class_var(self):
-        # Use class-level variables
-        result = VariableContainer.class_var + " used in method"
-        other_result = VariableContainer.reassignable_class_var
-        return result, other_result
-'''
-
-# Expected deleted lines for TypeScript DemoClass
-EXPECTED_DELETED_DEMOCLASS_TYPESCRIPT = """export class DemoClass {
-    value: number;
-    constructor(value: number) {
-        this.value = value;
-    }
-    printValue() {
-        console.log(this.value);
-    }
-}"""
-
 
 class DeleteSymbolTest(EditingTest):
-    def __init__(self, language: Language, rel_path: str, deleted_symbol: str, expected_deleted_lines: str):
+    def __init__(self, language: Language, rel_path: str, deleted_symbol: str):
         super().__init__(language, rel_path)
-        self.expected_deleted_lines = expected_deleted_lines
         self.deleted_symbol = deleted_symbol
         self.rel_path = rel_path
 
     def _apply_edit(self, symbol_manager: SymbolManager) -> None:
         symbol_manager.delete_symbol(self.deleted_symbol, self.rel_path)
 
-    def _test_diff(self, code_diff: CodeDiff) -> None:
-        assert code_diff.original_content != code_diff.modified_content
-        actual_deleted_lines = [line.strip() for _, line in code_diff.deleted_lines if line.strip()]
-        normalized_expected_deleted_lines = [line.strip() for line in self.expected_deleted_lines.splitlines() if line.strip()]
-        assert actual_deleted_lines == normalized_expected_deleted_lines
-
 
 @pytest.mark.parametrize(
     "test_case",
@@ -142,18 +89,16 @@ class DeleteSymbolTest(EditingTest):
             Language.PYTHON,
             PYTHON_TEST_REL_FILE_PATH,
             "VariableContainer",
-            EXPECTED_DELETED_VARIABLE_CONTAINER_PYTHON,
         ),
         DeleteSymbolTest(
             Language.TYPESCRIPT,
             TYPESCRIPT_TEST_FILE,
             "DemoClass",
-            EXPECTED_DELETED_DEMOCLASS_TYPESCRIPT,
         ),
     ],
 )
-def test_delete_symbol(test_case):
-    test_case.run_test()
+def test_delete_symbol(test_case, snapshot):
+    test_case.run_test(content_after_ground_truth=snapshot)
 
 
 NEW_PYTHON_FUNCTION = """def new_inserted_function():
@@ -168,47 +113,50 @@ NEW_PYTHON_VARIABLE = 'new_module_var = "Inserted after typed_module_var"'
 
 NEW_TYPESCRIPT_FUNCTION_AFTER = """function newFunctionAfterClass(): void {
     console.log("This function is after DemoClass.");
-}"""
+# }"""
 
 
-class InsertBeforeSymbolTest(EditingTest):
+class InsertInRelToSymbolTest(EditingTest):
     def __init__(self, language: Language, rel_path: str, symbol_name: str, new_content: str):
         super().__init__(language, rel_path)
         self.symbol_name = symbol_name
         self.new_content = new_content
+        self.mode: Literal["before", "after"] | None = None
+
+    def set_mode(self, mode: Literal["before", "after"]):
+        self.mode = mode
 
     def _apply_edit(self, symbol_manager: SymbolManager) -> None:
-        symbol_manager.insert_before_symbol(self.symbol_name, self.rel_path, self.new_content)
-
-    def _test_diff(self, code_diff: CodeDiff) -> None:
-        assert code_diff.original_content != code_diff.modified_content
-        actual_added_lines = [line.strip() for _, line in code_diff.added_lines if line.strip()]
-        normalized_expected_added_lines = [line.strip() for line in self.new_content.splitlines() if line.strip()]
-        assert actual_added_lines == normalized_expected_added_lines
+        assert self.mode is not None
+        if self.mode == "before":
+            symbol_manager.insert_before_symbol(self.symbol_name, self.rel_path, self.new_content)
+        elif self.mode == "after":
+            symbol_manager.insert_after_symbol(self.symbol_name, self.rel_path, self.new_content)
 
 
+@pytest.mark.parametrize("mode", ["before", "after"])
 @pytest.mark.parametrize(
     "test_case",
     [
-        InsertBeforeSymbolTest(
+        InsertInRelToSymbolTest(
             Language.PYTHON,
             PYTHON_TEST_REL_FILE_PATH,
             "typed_module_var",
             NEW_PYTHON_VARIABLE,
         ),
-        InsertBeforeSymbolTest(
+        InsertInRelToSymbolTest(
             Language.PYTHON,
             PYTHON_TEST_REL_FILE_PATH,
             "use_module_variables",
             NEW_PYTHON_FUNCTION,
         ),
-        InsertBeforeSymbolTest(
+        InsertInRelToSymbolTest(
             Language.TYPESCRIPT,
             TYPESCRIPT_TEST_FILE,
             "DemoClass",
             NEW_TYPESCRIPT_FUNCTION_AFTER,
         ),
-        InsertBeforeSymbolTest(
+        InsertInRelToSymbolTest(
             Language.TYPESCRIPT,
             TYPESCRIPT_TEST_FILE,
             "helperFunction",
@@ -216,8 +164,9 @@ class InsertBeforeSymbolTest(EditingTest):
         ),
     ],
 )
-def test_insert_before_symbol(test_case):
-    test_case.run_test()
+def test_insert_in_rel_to_symbol(test_case, mode, snapshot):
+    test_case.set_mode(mode)
+    test_case.run_test(content_after_ground_truth=snapshot)
 
 
 PYTHON_REPLACED_BODY = """        # This body has been replaced
@@ -229,19 +178,9 @@ TYPESCRIPT_REPLACED_BODY = """        // This body has been replaced
         console.warn("New value: " + this.value);
 """
 
-EXPECTED_ORIGINAL_MODIFY_INSTANCE_VAR_PYTHON = """    def modify_instance_var(self):
-        # Reassign instance variable
-        self.instance_var = "Modified instance value"
-        self.reassignable_instance_var = 200  # Reassigned"""
-
-# For single line original content, direct list is fine
-EXPECTED_ORIGINAL_PRINTVALUE_TYPESCRIPT = ["    printValue() {", "        console.log(this.value);", "    }"]
-
 
 class ReplaceBodyTest(EditingTest):
-    def __init__(
-        self, language: Language, rel_path: str, symbol_name: str, new_body: str, expected_removed_lines: str, expected_added_lines: str
-    ):
+    def __init__(self, language: Language, rel_path: str, symbol_name: str, new_body: str):
         super().__init__(language, rel_path)
         self.symbol_name = symbol_name
         self.new_body = new_body
@@ -249,16 +188,23 @@ class ReplaceBodyTest(EditingTest):
     def _apply_edit(self, symbol_manager: SymbolManager) -> None:
         symbol_manager.replace_body(self.symbol_name, self.rel_path, self.new_body)
 
-    def _test_diff(self, code_diff: CodeDiff) -> None:
-        assert code_diff.original_content != code_diff.modified_content
-        # TODO test properly
-
 
 @pytest.mark.parametrize(
     "test_case",
     [
-        # TODO add tests
+        ReplaceBodyTest(
+            Language.PYTHON,
+            PYTHON_TEST_REL_FILE_PATH,
+            "VariableContainer/modify_instance_var",
+            PYTHON_REPLACED_BODY,
+        ),
+        ReplaceBodyTest(
+            Language.TYPESCRIPT,
+            TYPESCRIPT_TEST_FILE,
+            "DemoClass/printValue",
+            TYPESCRIPT_REPLACED_BODY,
+        ),
     ],
 )
-def test_replace_body(test_case: ReplaceBodyTest):
-    test_case.run_test()
+def test_replace_body(test_case: ReplaceBodyTest, snapshot):
+    test_case.run_test(content_after_ground_truth=snapshot)

uv.lock

@@ -1006,6 +1006,7 @@ dev = [
     { name = "poethepoet" },
     { name = "pytest" },
     { name = "ruff" },
+    { name = "syrupy" },
     { name = "toml-sort" },
     { name = "types-pyyaml" },
 ]
@@ -1043,6 +1044,7 @@ requires-dist = [
     { name = "ruff", marker = "extra == 'dev'", specifier = ">=0.0.285" },
     { name = "sensai-utils", specifier = ">=1.4.0" },
     { name = "sqlalchemy", marker = "extra == 'agno'", specifier = ">=2.0.40" },
+    { name = "syrupy", marker = "extra == 'dev'", specifier = ">=4.9.1" },
     { name = "toml-sort", marker = "extra == 'dev'", specifier = ">=0.24.2" },
     { name = "types-pyyaml", specifier = ">=6.0.12.20241230" },
     { name = "types-pyyaml", marker = "extra == 'dev'", specifier = ">=6.0.12.20241230" },
@@ -1136,6 +1138,18 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/a0/4b/528ccf7a982216885a1ff4908e886b8fb5f19862d1962f56a3fce2435a70/starlette-0.46.1-py3-none-any.whl", hash = "sha256:77c74ed9d2720138b25875133f3a2dae6d854af2ec37dceb56aef370c1d8a227", size = 71995, upload-time = "2025-03-08T10:55:32.662Z" },
 ]
 
+[[package]]
+name = "syrupy"
+version = "4.9.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "pytest" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/8c/f8/022d8704a3314f3e96dbd6bbd16ebe119ce30e35f41aabfa92345652fceb/syrupy-4.9.1.tar.gz", hash = "sha256:b7d0fcadad80a7d2f6c4c71917918e8ebe2483e8c703dfc8d49cdbb01081f9a4", size = 52492, upload-time = "2025-03-24T01:36:37.225Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/ec/9d/aef9ec5fd5a4ee2f6a96032c4eda5888c5c7cec65cef6b28c4fc37671d88/syrupy-4.9.1-py3-none-any.whl", hash = "sha256:b94cc12ed0e5e75b448255430af642516842a2374a46936dd2650cfb6dd20eda", size = 52214, upload-time = "2025-03-24T01:36:35.278Z" },
+]
+
 [[package]]
 name = "tokenize-rt"
 version = "6.1.0"