adapt test

Author: themisvaltinos

tests/lsp/test_document_highlight.py

@@ -106,5 +106,6 @@ def test_get_document_highlights_multiple_ctes():
     )
     highlights = get_document_highlights(lsp_context, test_uri, position)
 
+    # This should return the column usages as well
     assert highlights is not None
-    assert len(highlights) == len(inner_ranges)  # Should match all occurrences of inner CTE only
+    assert len(highlights) == 4

tests/lsp/test_rename_cte.py

@@ -119,7 +119,9 @@ def test_rename_cte():
     assert uri in workspace_edit.changes
 
     edits = workspace_edit.changes[uri]
-    assert len(edits) == 2  # Should have 2 edits: definition + usage
+
+    # Should have edited four occurences including column usages
+    assert len(edits) == 4
 
     # Verify that both ranges are being edited
     edit_ranges = [edit.range for edit in edits]
@@ -135,6 +137,33 @@ def test_rename_cte():
     # Verify that all edits have the new name
     assert all(edit.new_text == "new_marketing" for edit in edits)
 
+    # Apply the edits to verify the result
+    with open(sushi_customers_path, "r", encoding="utf-8") as file:
+        lines = file.readlines()
+
+    # Apply edits in reverse order to avoid offset issues
+    sorted_edits = sorted(
+        edits, key=lambda e: (e.range.start.line, e.range.start.character), reverse=True
+    )
+    for edit in sorted_edits:
+        line_idx = edit.range.start.line
+        start_char = edit.range.start.character
+        end_char = edit.range.end.character
+
+        line = lines[line_idx]
+        new_line = line[:start_char] + edit.new_text + line[end_char:]
+        lines[line_idx] = new_line
+
+    # Verify the edited content
+    edited_content = "".join(lines)
+    assert "new_marketing" in edited_content
+    assert "current_marketing" not in edited_content.replace("current_marketing_outer", "")
+    assert edited_content.count("new_marketing") == 4
+    assert (
+        "  SELECT new_marketing.* FROM new_marketing WHERE new_marketing.customer_id != 100\n"
+        in lines
+    )
+
 
 def test_rename_cte_outer():
     context = Context(paths=["examples/sushi"])