From 07ea1370dc5e4e3f965b6106487e4ed98d2aaa9c Mon Sep 17 00:00:00 2001
From: Dominik Jain <dominik.jain@outlook.com>
Date: Mon, 23 Jun 2025 00:31:07 +0200
Subject: [PATCH] Add tests for adding a class with 'LLM-provided'
 leading/trailing newlines

---
 .../__snapshots__/test_symbol_editing.ambr    | 208 ++++++++++++++++++
 test/serena/test_symbol_editing.py            |  40 +++-
 2 files changed, 246 insertions(+), 2 deletions(-)

diff --git a/test/serena/__snapshots__/test_symbol_editing.ambr b/test/serena/__snapshots__/test_symbol_editing.ambr
index 2b698a2..b1f12e9 100644
--- a/test/serena/__snapshots__/test_symbol_editing.ambr
+++ b/test/serena/__snapshots__/test_symbol_editing.ambr
@@ -584,6 +584,214 @@
   
   '''
 # ---
+# name: test_insert_python_class_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"
+  
+  class NewInsertedClass:
+      pass
+  
+  
+  # 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_python_class_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
+  class NewInsertedClass:
+      pass
+  
+  @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_case0]
   '''
   """
diff --git a/test/serena/test_symbol_editing.py b/test/serena/test_symbol_editing.py
index b034d7c..29976a5 100644
--- a/test/serena/test_symbol_editing.py
+++ b/test/serena/test_symbol_editing.py
@@ -134,6 +134,18 @@ def test_delete_symbol(test_case, snapshot):
 NEW_PYTHON_FUNCTION = """def new_inserted_function():
     print("This is a new function inserted before another.")"""
 
+NEW_PYTHON_CLASS_WITH_LEADING_NEWLINES = """
+
+class NewInsertedClass:
+    pass
+"""
+
+NEW_PYTHON_CLASS_WITH_TRAILING_NEWLINES = """class NewInsertedClass:
+    pass
+
+
+"""
+
 NEW_TYPESCRIPT_FUNCTION = """function newInsertedFunction(): void {
     console.log("This is a new function inserted before another.");
 }"""
@@ -147,11 +159,13 @@ NEW_TYPESCRIPT_FUNCTION_AFTER = """function newFunctionAfterClass(): void {
 
 
 class InsertInRelToSymbolTest(EditingTest):
-    def __init__(self, language: Language, rel_path: str, symbol_name: str, new_content: str):
+    def __init__(
+        self, language: Language, rel_path: str, symbol_name: str, new_content: str, mode: Literal["before", "after"] | None = None
+    ):
         super().__init__(language, rel_path)
         self.symbol_name = symbol_name
         self.new_content = new_content
-        self.mode: Literal["before", "after"] | None = None
+        self.mode: Literal["before", "after"] | None = mode
 
     def set_mode(self, mode: Literal["before", "after"]):
         self.mode = mode
@@ -211,6 +225,28 @@ def test_insert_in_rel_to_symbol(test_case: InsertInRelToSymbolTest, mode: Liter
     test_case.run_test(content_after_ground_truth=snapshot)
 
 
+@pytest.mark.python
+def test_insert_python_class_before(snapshot):
+    InsertInRelToSymbolTest(
+        Language.PYTHON,
+        PYTHON_TEST_REL_FILE_PATH,
+        "VariableDataclass",
+        NEW_PYTHON_CLASS_WITH_TRAILING_NEWLINES,
+        mode="before",
+    ).run_test(snapshot)
+
+
+@pytest.mark.python
+def test_insert_python_class_after(snapshot):
+    InsertInRelToSymbolTest(
+        Language.PYTHON,
+        PYTHON_TEST_REL_FILE_PATH,
+        "VariableDataclass",
+        NEW_PYTHON_CLASS_WITH_LEADING_NEWLINES,
+        mode="after",
+    ).run_test(snapshot)
+
+
 PYTHON_REPLACED_BODY = """def modify_instance_var(self):
         # This body has been replaced
         self.instance_var = "Replaced!"