From 42514f7478ae3bdfd187434e52e3ea8fa9e42f9d Mon Sep 17 00:00:00 2001 From: Themis Valtinos <73662635+themisvaltinos@users.noreply.github.com> Date: Wed, 11 Jun 2025 18:05:06 +0300 Subject: [PATCH 1/4] Feat(lsp): Add rename functionality for CTEs in vscode --- sqlmesh/lsp/main.py | 54 ++++++++ sqlmesh/lsp/rename.py | 137 ++++++++++++++++++++ tests/lsp/test_document_highlight.py | 110 ++++++++++++++++ tests/lsp/test_rename_cte.py | 183 +++++++++++++++++++++++++++ 4 files changed, 484 insertions(+) create mode 100644 sqlmesh/lsp/rename.py create mode 100644 tests/lsp/test_document_highlight.py create mode 100644 tests/lsp/test_rename_cte.py diff --git a/sqlmesh/lsp/main.py b/sqlmesh/lsp/main.py index 2a3e18c6ea..7462698f7d 100755 --- a/sqlmesh/lsp/main.py +++ b/sqlmesh/lsp/main.py @@ -52,6 +52,7 @@ get_references, get_all_references, ) +from sqlmesh.lsp.rename import prepare_rename, rename_symbol, get_document_highlights from sqlmesh.lsp.uri import URI from web.server.api.endpoints.lineage import column_lineage, model_lineage from web.server.api.endpoints.models import get_models @@ -435,6 +436,59 @@ def find_references( ls.show_message(f"Error getting locations: {e}", types.MessageType.Error) return None + @self.server.feature(types.TEXT_DOCUMENT_PREPARE_RENAME) + def prepare_rename_handler( + ls: LanguageServer, params: types.PrepareRenameParams + ) -> t.Optional[types.PrepareRenameResult]: + """Prepare for rename operation by checking if the symbol can be renamed.""" + try: + uri = URI(params.text_document.uri) + self._ensure_context_for_document(uri) + if self.lsp_context is None: + raise RuntimeError(f"No context found for document: {uri}") + + result = prepare_rename(self.lsp_context, uri, params.position) + return result + except Exception as e: + ls.log_trace(f"Error preparing rename: {e}") + return None + + @self.server.feature(types.TEXT_DOCUMENT_RENAME) + def rename_handler( + ls: LanguageServer, params: types.RenameParams + ) -> t.Optional[types.WorkspaceEdit]: + """Perform rename operation on the symbol at the given position.""" + try: + uri = URI(params.text_document.uri) + self._ensure_context_for_document(uri) + if self.lsp_context is None: + raise RuntimeError(f"No context found for document: {uri}") + + workspace_edit = rename_symbol( + self.lsp_context, uri, params.position, params.new_name + ) + return workspace_edit + except Exception as e: + ls.show_message(f"Error performing rename: {e}", types.MessageType.Error) + return None + + @self.server.feature(types.TEXT_DOCUMENT_DOCUMENT_HIGHLIGHT) + def document_highlight_handler( + ls: LanguageServer, params: types.DocumentHighlightParams + ) -> t.Optional[t.List[types.DocumentHighlight]]: + """Highlight all occurrences of the symbol at the given position.""" + try: + uri = URI(params.text_document.uri) + self._ensure_context_for_document(uri) + if self.lsp_context is None: + raise RuntimeError(f"No context found for document: {uri}") + + highlights = get_document_highlights(self.lsp_context, uri, params.position) + return highlights + except Exception as e: + ls.log_trace(f"Error getting document highlights: {e}") + return None + @self.server.feature(types.TEXT_DOCUMENT_DIAGNOSTIC) def diagnostic( ls: LanguageServer, params: types.DocumentDiagnosticParams diff --git a/sqlmesh/lsp/rename.py b/sqlmesh/lsp/rename.py new file mode 100644 index 0000000000..0dbe2594ea --- /dev/null +++ b/sqlmesh/lsp/rename.py @@ -0,0 +1,137 @@ +import typing as t +from lsprotocol.types import ( + Position, + TextEdit, + WorkspaceEdit, + PrepareRenameResult_Type1, + DocumentHighlight, + DocumentHighlightKind, +) + +from sqlmesh.lsp.context import LSPContext +from sqlmesh.lsp.reference import ( + _position_within_range, + get_cte_references, + LSPCteReference, +) +from sqlmesh.lsp.uri import URI + + +def prepare_rename( + lsp_context: LSPContext, document_uri: URI, position: Position +) -> t.Optional[PrepareRenameResult_Type1]: + """ + Prepare for rename operation by checking if the symbol at the position can be renamed. + + Args: + lsp_context: The LSP context + document_uri: The URI of the document + position: The position in the document + + Returns: + PrepareRenameResult if the symbol can be renamed, None otherwise + """ + # Check if there's a CTE at this position + cte_references = get_cte_references(lsp_context, document_uri, position) + if cte_references: + # Find the target CTE definition to get its range + target_range = None + for ref in cte_references: + # Check if cursor is on a CTE usage + if _position_within_range(position, ref.range): + target_range = ref.target_range + break + # Check if cursor is on the CTE definition + elif _position_within_range(position, ref.target_range): + target_range = ref.target_range + break + if target_range: + return PrepareRenameResult_Type1(range=target_range, placeholder="cte_name") + + # For now, only CTEs are supported + return None + + +def rename_symbol( + lsp_context: LSPContext, document_uri: URI, position: Position, new_name: str +) -> t.Optional[WorkspaceEdit]: + """ + Perform rename operation on the symbol at the given position. + + Args: + lsp_context: The LSP context + document_uri: The URI of the document + position: The position in the document + new_name: The new name for the symbol + + Returns: + WorkspaceEdit with the changes, or None if no symbol to rename + """ + # Check if there's a CTE at this position + cte_references = get_cte_references(lsp_context, document_uri, position) + if cte_references: + return _rename_cte(cte_references, new_name) + + # For now, only CTEs are supported + return None + + +def _rename_cte(cte_references: t.List[LSPCteReference], new_name: str) -> WorkspaceEdit: + """ + Create a WorkspaceEdit for renaming a CTE. + + Args: + cte_references: List of CTE references (definition and usages) + new_name: The new name for the CTE + + Returns: + WorkspaceEdit with the text edits for renaming the CTE + """ + changes: t.Dict[str, t.List[TextEdit]] = {} + + for ref in cte_references: + uri = ref.uri + if uri not in changes: + changes[uri] = [] + + # Create a text edit for this reference + text_edit = TextEdit(range=ref.range, new_text=new_name) + changes[uri].append(text_edit) + + return WorkspaceEdit(changes=changes) + + +def get_document_highlights( + lsp_context: LSPContext, document_uri: URI, position: Position +) -> t.Optional[t.List[DocumentHighlight]]: + """ + Get document highlights for all occurrences of the symbol at the given position. + + This function finds all occurrences of a symbol (CTE) within the current document + and returns them as DocumentHighlight objects for "Change All Occurrences" feature. + + Args: + lsp_context: The LSP context + document_uri: The URI of the document + position: The position in the document to find highlights for + + Returns: + List of DocumentHighlight objects or None if no symbol found + """ + # Check if there's a CTE at this position + cte_references = get_cte_references(lsp_context, document_uri, position) + if cte_references: + highlights = [] + for ref in cte_references: + # Determine the highlight kind based on whether it's a definition or usage + kind = ( + DocumentHighlightKind.Write + if ref.range == ref.target_range + else DocumentHighlightKind.Read + ) + + highlights.append(DocumentHighlight(range=ref.range, kind=kind)) + return highlights + + # For now, only CTEs are supported + return None diff --git a/tests/lsp/test_document_highlight.py b/tests/lsp/test_document_highlight.py new file mode 100644 index 0000000000..12688911e8 --- /dev/null +++ b/tests/lsp/test_document_highlight.py @@ -0,0 +1,110 @@ +from lsprotocol.types import Position, DocumentHighlightKind + +from sqlmesh.core.context import Context +from sqlmesh.lsp.context import LSPContext, ModelTarget +from sqlmesh.lsp.rename import get_document_highlights +from sqlmesh.lsp.uri import URI +from tests.lsp.test_reference_cte import find_ranges_from_regex + + +def test_get_document_highlights_cte(): + context = Context(paths=["examples/sushi"]) + lsp_context = LSPContext(context) + + # Use the existing customers.sql model which has CTEs + sushi_customers_path = next( + path + for path, info in lsp_context.map.items() + if isinstance(info, ModelTarget) and "sushi.customers" in info.names + ) + + with open(sushi_customers_path, "r", encoding="utf-8") as file: + read_file = file.readlines() + + test_uri = URI.from_path(sushi_customers_path) + + # Find the ranges for "current_marketing" CTE (not outer one) + ranges = find_ranges_from_regex(read_file, r"current_marketing(?!_outer)") + assert len(ranges) >= 2 # Should have definition + usage + + # Test highlighting CTE definition - position on "current_marketing" definition + position = Position(line=ranges[0].start.line, character=ranges[0].start.character + 4) + highlights = get_document_highlights(lsp_context, test_uri, position) + + assert highlights is not None + assert len(highlights) >= 2 # Definition + at least 1 usage + + # Check that we have both definition (Write) and usage (Read) highlights + highlight_kinds = [h.kind for h in highlights] + assert DocumentHighlightKind.Write in highlight_kinds # CTE definition + assert DocumentHighlightKind.Read in highlight_kinds # CTE usage + + # Test highlighting CTE usage - position on "current_marketing" usage + position = Position(line=ranges[1].start.line, character=ranges[1].start.character + 4) + highlights = get_document_highlights(lsp_context, test_uri, position) + + assert highlights is not None + assert len(highlights) >= 2 # Should find the same references + + +def test_get_document_highlights_no_symbol(): + context = Context(paths=["examples/sushi"]) + lsp_context = LSPContext(context) + + # Use the existing customers.sql model + sushi_customers_path = next( + path + for path, info in lsp_context.map.items() + if isinstance(info, ModelTarget) and "sushi.customers" in info.names + ) + + test_uri = URI.from_path(sushi_customers_path) + + # Test position not on any CTE symbol - just on a random keyword + position = Position(line=5, character=5) + highlights = get_document_highlights(lsp_context, test_uri, position) + + assert highlights is None + + +def test_get_document_highlights_multiple_ctes(): + context = Context(paths=["examples/sushi"]) + lsp_context = LSPContext(context) + + # Use the existing customers.sql model which has both outer and inner CTEs + sushi_customers_path = next( + path + for path, info in lsp_context.map.items() + if isinstance(info, ModelTarget) and "sushi.customers" in info.names + ) + + with open(sushi_customers_path, "r", encoding="utf-8") as file: + read_file = file.readlines() + + test_uri = URI.from_path(sushi_customers_path) + + # Test the outer CTE - "current_marketing_outer" + outer_ranges = find_ranges_from_regex(read_file, r"current_marketing_outer") + assert len(outer_ranges) >= 2 # Should have definition + usage + + # Test highlighting outer CTE - should only highlight that CTE + position = Position( + line=outer_ranges[0].start.line, character=outer_ranges[0].start.character + 4 + ) + highlights = get_document_highlights(lsp_context, test_uri, position) + + assert highlights is not None + assert len(highlights) == len(outer_ranges) # Should match all occurrences of outer CTE + + # Test the inner CTE - "current_marketing" (not outer) + inner_ranges = find_ranges_from_regex(read_file, r"current_marketing(?!_outer)") + assert len(inner_ranges) >= 2 # Should have definition + usage + + # Test highlighting inner CTE - should only highlight that CTE, not the outer one + position = Position( + line=inner_ranges[0].start.line, character=inner_ranges[0].start.character + 4 + ) + highlights = get_document_highlights(lsp_context, test_uri, position) + + assert highlights is not None + assert len(highlights) == len(inner_ranges) # Should match all occurrences of inner CTE only diff --git a/tests/lsp/test_rename_cte.py b/tests/lsp/test_rename_cte.py new file mode 100644 index 0000000000..b1927e29fb --- /dev/null +++ b/tests/lsp/test_rename_cte.py @@ -0,0 +1,183 @@ +from lsprotocol.types import Position +from sqlmesh.core.context import Context +from sqlmesh.lsp.context import LSPContext, ModelTarget +from sqlmesh.lsp.rename import prepare_rename, rename_symbol +from sqlmesh.lsp.uri import URI +from tests.lsp.test_reference_cte import find_ranges_from_regex + + +def test_prepare_rename_cte(): + context = Context(paths=["examples/sushi"]) + lsp_context = LSPContext(context) + + sushi_customers_path = next( + path + for path, info in lsp_context.map.items() + if isinstance(info, ModelTarget) and "sushi.customers" in info.names + ) + + with open(sushi_customers_path, "r", encoding="utf-8") as file: + read_file = file.readlines() + + # Test clicking on CTE definition for "current_marketing" + ranges = find_ranges_from_regex(read_file, r"current_marketing(?!_outer)") + assert len(ranges) == 2 + + # Click on the CTE definition + position = Position(line=ranges[0].start.line, character=ranges[0].start.character + 4) + result = prepare_rename(lsp_context, URI.from_path(sushi_customers_path), position) + + assert result is not None + assert result.placeholder == "cte_name" + assert result.range == ranges[0] # Should return the definition range + + # Test clicking on CTE usage + position = Position(line=ranges[1].start.line, character=ranges[1].start.character + 4) + result = prepare_rename(lsp_context, URI.from_path(sushi_customers_path), position) + + assert result is not None + assert result.placeholder == "cte_name" + assert result.range == ranges[0] # Should still return the definition range + + +def test_prepare_rename_cte_outer(): + context = Context(paths=["examples/sushi"]) + lsp_context = LSPContext(context) + + sushi_customers_path = next( + path + for path, info in lsp_context.map.items() + if isinstance(info, ModelTarget) and "sushi.customers" in info.names + ) + + with open(sushi_customers_path, "r", encoding="utf-8") as file: + read_file = file.readlines() + + # Test clicking on CTE definition for "current_marketing_outer" + ranges = find_ranges_from_regex(read_file, r"current_marketing_outer") + assert len(ranges) == 2 + + # Click on the CTE definition + position = Position(line=ranges[0].start.line, character=ranges[0].start.character + 4) + result = prepare_rename(lsp_context, URI.from_path(sushi_customers_path), position) + + assert result is not None + assert result.placeholder == "cte_name" + assert result.range == ranges[0] # Should return the definition range + + +def test_prepare_rename_non_cte(): + context = Context(paths=["examples/sushi"]) + lsp_context = LSPContext(context) + + sushi_customers_path = next( + path + for path, info in lsp_context.map.items() + if isinstance(info, ModelTarget) and "sushi.customers" in info.names + ) + + with open(sushi_customers_path, "r", encoding="utf-8") as file: + read_file = file.readlines() + + # Click on a regular table reference (not a CTE) + ranges = find_ranges_from_regex(read_file, r"sushi\.orders") + assert len(ranges) >= 1 + + position = Position(line=ranges[0].start.line, character=ranges[0].start.character + 4) + result = prepare_rename(lsp_context, URI.from_path(sushi_customers_path), position) + + assert result is None + + +def test_rename_cte(): + context = Context(paths=["examples/sushi"]) + lsp_context = LSPContext(context) + + sushi_customers_path = next( + path + for path, info in lsp_context.map.items() + if isinstance(info, ModelTarget) and "sushi.customers" in info.names + ) + + with open(sushi_customers_path, "r", encoding="utf-8") as file: + read_file = file.readlines() + + # Test renaming "current_marketing" to "new_marketing" + ranges = find_ranges_from_regex(read_file, r"current_marketing(?!_outer)") + assert len(ranges) == 2 + + # Click on the CTE definition + position = Position(line=ranges[0].start.line, character=ranges[0].start.character + 4) + workspace_edit = rename_symbol( + lsp_context, URI.from_path(sushi_customers_path), position, "new_marketing" + ) + + assert workspace_edit is not None + assert workspace_edit.changes is not None + + uri = URI.from_path(sushi_customers_path).value + assert uri in workspace_edit.changes + + edits = workspace_edit.changes[uri] + assert len(edits) == 2 # Should have 2 edits: definition + usage + + # Verify that both ranges are being edited + edit_ranges = [edit.range for edit in edits] + for expected_range in ranges: + assert any( + edit_range.start.line == expected_range.start.line + and edit_range.start.character == expected_range.start.character + for edit_range in edit_ranges + ), ( + f"Expected to find edit at line {expected_range.start.line}, char {expected_range.start.character}" + ) + + # Verify that all edits have the new name + assert all(edit.new_text == "new_marketing" for edit in edits) + + +def test_rename_cte_outer(): + context = Context(paths=["examples/sushi"]) + lsp_context = LSPContext(context) + + sushi_customers_path = next( + path + for path, info in lsp_context.map.items() + if isinstance(info, ModelTarget) and "sushi.customers" in info.names + ) + + with open(sushi_customers_path, "r", encoding="utf-8") as file: + read_file = file.readlines() + + # Test renaming "current_marketing_outer" to "new_marketing_outer" + ranges = find_ranges_from_regex(read_file, r"current_marketing_outer") + assert len(ranges) == 2 + + # Click on the CTE usage + position = Position(line=ranges[1].start.line, character=ranges[1].start.character + 4) + workspace_edit = rename_symbol( + lsp_context, URI.from_path(sushi_customers_path), position, "new_marketing_outer" + ) + + assert workspace_edit is not None + assert workspace_edit.changes is not None + + uri = URI.from_path(sushi_customers_path).value + assert uri in workspace_edit.changes + + edits = workspace_edit.changes[uri] + assert len(edits) == 2 # Should have 2 edits: definition + usage + + # Verify that both ranges are being edited + edit_ranges = [edit.range for edit in edits] + for expected_range in ranges: + assert any( + edit_range.start.line == expected_range.start.line + and edit_range.start.character == expected_range.start.character + for edit_range in edit_ranges + ), ( + f"Expected to find edit at line {expected_range.start.line}, char {expected_range.start.character}" + ) + + # Verify that all edits have the new name + assert all(edit.new_text == "new_marketing_outer" for edit in edits) From f13ce21afd6dca51236e44151fd058cce9354fc6 Mon Sep 17 00:00:00 2001 From: Themis Valtinos <73662635+themisvaltinos@users.noreply.github.com> Date: Fri, 20 Jun 2025 14:25:48 +0300 Subject: [PATCH 2/4] adapt test --- tests/lsp/test_document_highlight.py | 3 ++- tests/lsp/test_rename_cte.py | 31 +++++++++++++++++++++++++++- 2 files changed, 32 insertions(+), 2 deletions(-) diff --git a/tests/lsp/test_document_highlight.py b/tests/lsp/test_document_highlight.py index 12688911e8..e6ce0ae7ec 100644 --- a/tests/lsp/test_document_highlight.py +++ b/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 diff --git a/tests/lsp/test_rename_cte.py b/tests/lsp/test_rename_cte.py index b1927e29fb..4ca1002c2e 100644 --- a/tests/lsp/test_rename_cte.py +++ b/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"]) From bb548860f3def3dac61e6898382238e040329005 Mon Sep 17 00:00:00 2001 From: Themis Valtinos <73662635+themisvaltinos@users.noreply.github.com> Date: Fri, 20 Jun 2025 15:31:35 +0300 Subject: [PATCH 3/4] add e2e test --- vscode/extension/tests/rename_cte.spec.ts | 187 ++++++++++++++++++++++ 1 file changed, 187 insertions(+) create mode 100644 vscode/extension/tests/rename_cte.spec.ts diff --git a/vscode/extension/tests/rename_cte.spec.ts b/vscode/extension/tests/rename_cte.spec.ts new file mode 100644 index 0000000000..9d68e63f6c --- /dev/null +++ b/vscode/extension/tests/rename_cte.spec.ts @@ -0,0 +1,187 @@ +import { test, expect } from '@playwright/test' +import path from 'path' +import fs from 'fs-extra' +import os from 'os' +import { startVSCode, SUSHI_SOURCE_PATH } from './utils' + +// Keyboard shortcuts +const RENAME_KEY = 'F2' +const FIND_ALL_REFERENCES_KEY = + process.platform === 'darwin' ? 'Alt+Shift+F12' : 'Ctrl+Shift+F12' + +test.describe('CTE Rename', () => { + let tempDir: string + let window: any + let close: () => Promise + + test.beforeEach(async () => { + tempDir = await fs.mkdtemp(path.join(os.tmpdir(), 'vscode-test-sushi-')) + await fs.copy(SUSHI_SOURCE_PATH, tempDir) + const vscode = await startVSCode(tempDir) + window = vscode.window + close = vscode.close + + // Navigate to customers.sql which contains CTEs + await window.waitForSelector('text=models') + await window + .getByRole('treeitem', { name: 'models', exact: true }) + .locator('a') + .click() + await window + .getByRole('treeitem', { name: 'customers.sql', exact: true }) + .locator('a') + .click() + await window.waitForSelector('text=grain') + await window.waitForSelector('text=Loaded SQLMesh Context') + }) + + test.afterEach(async () => { + await close() + fs.removeSync(tempDir) + }) + + test('Rename CTE from definition', async () => { + // Click on the inner CTE definition "current_marketing" (not the outer one) + await window.locator('text=WITH current_marketing AS').click({ + position: { x: 100, y: 5 }, + }) + + // Press F2 to trigger rename + await window.keyboard.press(RENAME_KEY) + await expect(window.locator('text=Rename')).toBeVisible() + const renameInput = window.locator('input:focus') + await expect(renameInput).toBeVisible() + + // Type new name and confirm + await window.keyboard.type('new_marketing') + await window.keyboard.press('Enter') + await window.waitForTimeout(1000) + + // Verify the rename was applied + await expect(window.locator('text=WITH new_marketing AS')).toBeVisible() + }) + + test('Rename CTE from usage', async () => { + // Click on CTE usage in FROM clause + await window.locator('text=FROM current_marketing_outer').click({ + position: { x: 80, y: 5 }, + }) + + // Press F2 to trigger rename + await window.keyboard.press(RENAME_KEY) + + // Wait for rename input to appear + await expect(window.locator('text=Rename')).toBeVisible() + const renameInput = window.locator('input:focus') + await expect(renameInput).toBeVisible() + + // Type new name + await window.keyboard.type('updated_marketing_out') + + // Confirm rename + await window.keyboard.press('Enter') + await window.waitForTimeout(1000) + + // Verify both definition and usage were renamed + await expect( + window.locator('text=WITH updated_marketing_out AS'), + ).toBeVisible() + await expect( + window.locator('text=FROM updated_marketing_out'), + ).toBeVisible() + }) + + test('Cancel CTE rename', async () => { + // Click on the CTE to rename + await window.locator('text=current_marketing_outer').first().click() + + // Press F2 to trigger rename + await window.keyboard.press(RENAME_KEY) + + // Wait for rename input to appear + await expect(window.locator('text=Rename')).toBeVisible() + const renameInput = window.locator('input:focus') + await expect(renameInput).toBeVisible() + + // Type new name but cancel + await window.keyboard.type('cancelled_name') + await window.keyboard.press('Escape') + + // Wait for UI to update + await window.waitForTimeout(500) + + // Verify CTE name was NOT changed + await expect( + window.locator('text=current_marketing_outer').first(), + ).toBeVisible() + await expect(window.locator('text=cancelled_name')).not.toBeVisible() + }) + + test('Rename CTE updates all references', async () => { + // Click on the CTE definition + await window.locator('text=WITH current_marketing AS').click({ + position: { x: 100, y: 5 }, + }) + + // Press F2 to trigger rename + await window.keyboard.press(RENAME_KEY) + // Wait for rename input to appear + await expect(window.locator('text=Rename')).toBeVisible() + const renameInput = window.locator('input:focus') + await expect(renameInput).toBeVisible() + + // Type new name and confirm + await window.keyboard.type('renamed_cte') + await window.keyboard.press('Enter') + + // Click on the renamed CTE + await window.locator('text=WITH renamed_cte AS').click({ + position: { x: 100, y: 5 }, + }) + + // Find all references using keyboard shortcut + await window.keyboard.press(FIND_ALL_REFERENCES_KEY) + + // Verify references panel shows all occurrences + await window.waitForSelector('text=References') + await expect(window.locator('text=customers.sql').first()).toBeVisible() + await window.waitForSelector('text=WITH renamed_cte AS') + await window.waitForSelector('text=renamed_cte.*') + await window.waitForSelector('text=FROM renamed_cte') + await window.waitForSelector('text=renamed_cte.customer_id != 100') + }) + + test('Rename CTE with preview', async () => { + // Click on the CTE to rename + await window.locator('text=WITH current_marketing AS').click({ + position: { x: 100, y: 5 }, + }) + + // Press F2 to trigger rename + await window.keyboard.press(RENAME_KEY) + await expect(window.locator('text=Rename')).toBeVisible() + const renameInput = window.locator('input:focus') + await expect(renameInput).toBeVisible() + + // Type new name + await window.keyboard.type('preview_marketing') + + // Press Cmd+Enter (Meta+Enter) to preview changes + await window.keyboard.press('Meta+Enter') + + // Verify preview UI is showing + await expect(window.locator('text=Refactor Preview').first()).toBeVisible() + await expect(window.locator('text=Apply').first()).toBeVisible() + await expect(window.locator('text=Discard').first()).toBeVisible() + + // Verify the preview shows both old and new names + await expect(window.locator('text=current_marketing').first()).toBeVisible() + await expect(window.locator('text=preview_marketing').first()).toBeVisible() + + // Apply the changes + await window.locator('text=Apply').click() + + // Verify the rename was applied + await expect(window.locator('text=WITH preview_marketing AS')).toBeVisible() + }) +}) From 813aa614976100a0a43b868c606d6304e1ccaa3a Mon Sep 17 00:00:00 2001 From: Themis Valtinos <73662635+themisvaltinos@users.noreply.github.com> Date: Mon, 23 Jun 2025 13:24:37 +0300 Subject: [PATCH 4/4] address comments --- vscode/extension/tests/rename_cte.spec.ts | 355 ++++++++++++---------- 1 file changed, 195 insertions(+), 160 deletions(-) diff --git a/vscode/extension/tests/rename_cte.spec.ts b/vscode/extension/tests/rename_cte.spec.ts index 9d68e63f6c..e1b5da6a7e 100644 --- a/vscode/extension/tests/rename_cte.spec.ts +++ b/vscode/extension/tests/rename_cte.spec.ts @@ -9,179 +9,214 @@ const RENAME_KEY = 'F2' const FIND_ALL_REFERENCES_KEY = process.platform === 'darwin' ? 'Alt+Shift+F12' : 'Ctrl+Shift+F12' -test.describe('CTE Rename', () => { - let tempDir: string - let window: any - let close: () => Promise - - test.beforeEach(async () => { - tempDir = await fs.mkdtemp(path.join(os.tmpdir(), 'vscode-test-sushi-')) - await fs.copy(SUSHI_SOURCE_PATH, tempDir) - const vscode = await startVSCode(tempDir) - window = vscode.window - close = vscode.close - - // Navigate to customers.sql which contains CTEs - await window.waitForSelector('text=models') - await window - .getByRole('treeitem', { name: 'models', exact: true }) - .locator('a') - .click() - await window - .getByRole('treeitem', { name: 'customers.sql', exact: true }) - .locator('a') - .click() - await window.waitForSelector('text=grain') - await window.waitForSelector('text=Loaded SQLMesh Context') - }) - - test.afterEach(async () => { - await close() - fs.removeSync(tempDir) - }) +// Helper function to set up a test environment +async function setupTestEnvironment() { + const tempDir = await fs.mkdtemp(path.join(os.tmpdir(), 'vscode-test-sushi-')) + await fs.copy(SUSHI_SOURCE_PATH, tempDir) + const { window, close } = await startVSCode(tempDir) + + // Navigate to customers.sql which contains CTEs + await window.waitForSelector('text=models') + await window + .getByRole('treeitem', { name: 'models', exact: true }) + .locator('a') + .click() + await window + .getByRole('treeitem', { name: 'customers.sql', exact: true }) + .locator('a') + .click() + await window.waitForSelector('text=grain') + await window.waitForSelector('text=Loaded SQLMesh Context') + + return { window, close, tempDir } +} +test.describe('CTE Rename', () => { test('Rename CTE from definition', async () => { - // Click on the inner CTE definition "current_marketing" (not the outer one) - await window.locator('text=WITH current_marketing AS').click({ - position: { x: 100, y: 5 }, - }) - - // Press F2 to trigger rename - await window.keyboard.press(RENAME_KEY) - await expect(window.locator('text=Rename')).toBeVisible() - const renameInput = window.locator('input:focus') - await expect(renameInput).toBeVisible() - - // Type new name and confirm - await window.keyboard.type('new_marketing') - await window.keyboard.press('Enter') - await window.waitForTimeout(1000) - - // Verify the rename was applied - await expect(window.locator('text=WITH new_marketing AS')).toBeVisible() + const { window, close, tempDir } = await setupTestEnvironment() + + try { + // Click on the inner CTE definition "current_marketing" (not the outer one) + await window.locator('text=WITH current_marketing AS').click({ + position: { x: 100, y: 5 }, + }) + + // Press F2 to trigger rename + await window.keyboard.press(RENAME_KEY) + await expect(window.locator('text=Rename')).toBeVisible() + const renameInput = window.locator('input:focus') + await expect(renameInput).toBeVisible() + + // Type new name and confirm + await window.keyboard.type('new_marketing') + await window.keyboard.press('Enter') + await window.waitForTimeout(1000) + + // Verify the rename was applied + await expect(window.locator('text=WITH new_marketing AS')).toBeVisible() + } finally { + await close() + await fs.remove(tempDir) + } }) test('Rename CTE from usage', async () => { - // Click on CTE usage in FROM clause - await window.locator('text=FROM current_marketing_outer').click({ - position: { x: 80, y: 5 }, - }) - - // Press F2 to trigger rename - await window.keyboard.press(RENAME_KEY) - - // Wait for rename input to appear - await expect(window.locator('text=Rename')).toBeVisible() - const renameInput = window.locator('input:focus') - await expect(renameInput).toBeVisible() - - // Type new name - await window.keyboard.type('updated_marketing_out') - - // Confirm rename - await window.keyboard.press('Enter') - await window.waitForTimeout(1000) - - // Verify both definition and usage were renamed - await expect( - window.locator('text=WITH updated_marketing_out AS'), - ).toBeVisible() - await expect( - window.locator('text=FROM updated_marketing_out'), - ).toBeVisible() + const { window, close, tempDir } = await setupTestEnvironment() + + try { + // Click on CTE usage in FROM clause + await window.locator('text=FROM current_marketing_outer').click({ + position: { x: 80, y: 5 }, + }) + + // Press F2 to trigger rename + await window.keyboard.press(RENAME_KEY) + + // Wait for rename input to appear + await expect(window.locator('text=Rename')).toBeVisible() + const renameInput = window.locator('input:focus') + await expect(renameInput).toBeVisible() + + // Type new name + await window.keyboard.type('updated_marketing_out') + + // Confirm rename + await window.keyboard.press('Enter') + await window.waitForTimeout(1000) + + // Verify both definition and usage were renamed + await expect( + window.locator('text=WITH updated_marketing_out AS'), + ).toBeVisible() + await expect( + window.locator('text=FROM updated_marketing_out'), + ).toBeVisible() + } finally { + await close() + await fs.remove(tempDir) + } }) test('Cancel CTE rename', async () => { - // Click on the CTE to rename - await window.locator('text=current_marketing_outer').first().click() - - // Press F2 to trigger rename - await window.keyboard.press(RENAME_KEY) - - // Wait for rename input to appear - await expect(window.locator('text=Rename')).toBeVisible() - const renameInput = window.locator('input:focus') - await expect(renameInput).toBeVisible() - - // Type new name but cancel - await window.keyboard.type('cancelled_name') - await window.keyboard.press('Escape') - - // Wait for UI to update - await window.waitForTimeout(500) - - // Verify CTE name was NOT changed - await expect( - window.locator('text=current_marketing_outer').first(), - ).toBeVisible() - await expect(window.locator('text=cancelled_name')).not.toBeVisible() + const { window, close, tempDir } = await setupTestEnvironment() + + try { + // Click on the CTE to rename + await window.locator('text=current_marketing_outer').first().click() + + // Press F2 to trigger rename + await window.keyboard.press(RENAME_KEY) + + // Wait for rename input to appear + await expect(window.locator('text=Rename')).toBeVisible() + const renameInput = window.locator('input:focus') + await expect(renameInput).toBeVisible() + + // Type new name but cancel + await window.keyboard.type('cancelled_name') + await window.keyboard.press('Escape') + + // Wait for UI to update + await window.waitForTimeout(500) + + // Verify CTE name was NOT changed + await expect( + window.locator('text=current_marketing_outer').first(), + ).toBeVisible() + await expect(window.locator('text=cancelled_name')).not.toBeVisible() + } finally { + await close() + await fs.remove(tempDir) + } }) test('Rename CTE updates all references', async () => { - // Click on the CTE definition - await window.locator('text=WITH current_marketing AS').click({ - position: { x: 100, y: 5 }, - }) - - // Press F2 to trigger rename - await window.keyboard.press(RENAME_KEY) - // Wait for rename input to appear - await expect(window.locator('text=Rename')).toBeVisible() - const renameInput = window.locator('input:focus') - await expect(renameInput).toBeVisible() - - // Type new name and confirm - await window.keyboard.type('renamed_cte') - await window.keyboard.press('Enter') - - // Click on the renamed CTE - await window.locator('text=WITH renamed_cte AS').click({ - position: { x: 100, y: 5 }, - }) - - // Find all references using keyboard shortcut - await window.keyboard.press(FIND_ALL_REFERENCES_KEY) - - // Verify references panel shows all occurrences - await window.waitForSelector('text=References') - await expect(window.locator('text=customers.sql').first()).toBeVisible() - await window.waitForSelector('text=WITH renamed_cte AS') - await window.waitForSelector('text=renamed_cte.*') - await window.waitForSelector('text=FROM renamed_cte') - await window.waitForSelector('text=renamed_cte.customer_id != 100') + const { window, close, tempDir } = await setupTestEnvironment() + + try { + // Click on the CTE definition + await window.locator('text=WITH current_marketing AS').click({ + position: { x: 100, y: 5 }, + }) + + // Press F2 to trigger rename + await window.keyboard.press(RENAME_KEY) + // Wait for rename input to appear + await expect(window.locator('text=Rename')).toBeVisible() + const renameInput = window.locator('input:focus') + await expect(renameInput).toBeVisible() + + // Type new name and confirm + await window.keyboard.type('renamed_cte') + await window.keyboard.press('Enter') + + // Click on the renamed CTE + await window.locator('text=WITH renamed_cte AS').click({ + position: { x: 100, y: 5 }, + }) + + // Find all references using keyboard shortcut + await window.keyboard.press(FIND_ALL_REFERENCES_KEY) + + // Verify references panel shows all occurrences + await window.waitForSelector('text=References') + await expect(window.locator('text=customers.sql').first()).toBeVisible() + await window.waitForSelector('text=WITH renamed_cte AS') + await window.waitForSelector('text=renamed_cte.*') + await window.waitForSelector('text=FROM renamed_cte') + await window.waitForSelector('text=renamed_cte.customer_id != 100') + } finally { + await close() + await fs.remove(tempDir) + } }) test('Rename CTE with preview', async () => { - // Click on the CTE to rename - await window.locator('text=WITH current_marketing AS').click({ - position: { x: 100, y: 5 }, - }) - - // Press F2 to trigger rename - await window.keyboard.press(RENAME_KEY) - await expect(window.locator('text=Rename')).toBeVisible() - const renameInput = window.locator('input:focus') - await expect(renameInput).toBeVisible() - - // Type new name - await window.keyboard.type('preview_marketing') - - // Press Cmd+Enter (Meta+Enter) to preview changes - await window.keyboard.press('Meta+Enter') - - // Verify preview UI is showing - await expect(window.locator('text=Refactor Preview').first()).toBeVisible() - await expect(window.locator('text=Apply').first()).toBeVisible() - await expect(window.locator('text=Discard').first()).toBeVisible() - - // Verify the preview shows both old and new names - await expect(window.locator('text=current_marketing').first()).toBeVisible() - await expect(window.locator('text=preview_marketing').first()).toBeVisible() - - // Apply the changes - await window.locator('text=Apply').click() - - // Verify the rename was applied - await expect(window.locator('text=WITH preview_marketing AS')).toBeVisible() + const { window, close, tempDir } = await setupTestEnvironment() + + try { + // Click on the CTE to rename + await window.locator('text=WITH current_marketing AS').click({ + position: { x: 100, y: 5 }, + }) + + // Press F2 to trigger rename + await window.keyboard.press(RENAME_KEY) + await expect(window.locator('text=Rename')).toBeVisible() + const renameInput = window.locator('input:focus') + await expect(renameInput).toBeVisible() + + // Type new name + await window.keyboard.type('preview_marketing') + + // Press Cmd+Enter (Meta+Enter) to preview changes + await window.keyboard.press('Meta+Enter') + + // Verify preview UI is showing + await expect( + window.locator('text=Refactor Preview').first(), + ).toBeVisible() + await expect(window.locator('text=Apply').first()).toBeVisible() + await expect(window.locator('text=Discard').first()).toBeVisible() + + // Verify the preview shows both old and new names + await expect( + window.locator('text=current_marketing').first(), + ).toBeVisible() + await expect( + window.locator('text=preview_marketing').first(), + ).toBeVisible() + + // Apply the changes + await window.locator('text=Apply').click() + + // Verify the rename was applied + await expect( + window.locator('text=WITH preview_marketing AS'), + ).toBeVisible() + } finally { + await close() + await fs.remove(tempDir) + } }) })