temp

[ci skip]

Author: benfdking

tests/lsp/test_reference_cte.py

@@ -1,12 +1,11 @@
-#!/usr/bin/env python3
-"""Test script for CTE go-to-definition functionality."""
-
 import re
 from sqlmesh.core.context import Context
 from sqlmesh.lsp.context import LSPContext, ModelTarget
 from sqlmesh.lsp.reference import get_references
 from sqlmesh.lsp.uri import URI
 from lsprotocol.types import Range, Position
+import typing as t
+
 
 def test_cte_parsing():
     context = Context(paths=["examples/sushi"])
@@ -22,29 +21,41 @@ def test_cte_parsing():
     with open(sushi_customers_path, "r", encoding="utf-8") as file:
         read_file = file.readlines()
 
-    references = get_references(lsp_context, URI.from_path(sushi_customers_path), Position(line=0, character=0))
+    # Find position of the cte reference 
+    ranges = find_ranges_from_regex(
+        read_file,
+        r"current_marketing"
+    )
+    assert len(ranges) == 2
+    # Middle of the second range
+    position = Position(line=ranges[1].start.line, character=ranges[1].start.character + 4)
+        
+    references = get_references(
+        lsp_context, URI.from_path(sushi_customers_path), position 
+    )
 
     assert len(references) == 1
     assert references[0].uri == URI.from_path(sushi_customers_path)
+    assert references[0].description is None
 
-    ranges = find_ranges_from_regex(read_file, r"WITH\s+current_marketing\s+AS\s+\(SELECT\s+customer_id,\s+status\s+FROM\s+sushi\.marketing\s+WHERE\s+valid_to\s+is\s+null\)")
+    ranges = find_ranges_from_regex(
+        read_file,
+        r"WITH\s+current_marketing\s+AS\s+\(SELECT\s+customer_id,\s+status\s+FROM\s+sushi\.marketing\s+WHERE\s+valid_to\s+is\s+null\)",
+    )
     assert len(ranges) == 1
     assert ranges[0].start == Position(line=1, character=0)
     assert ranges[0].end == Position(line=1, character=1)
 
-
-
+    assert references[0].range == ranges[0]
 
 
 def find_ranges_from_regex(read_file: t.List[str], regex: str) -> t.List[Range]:
     """Find all ranges in the read file that match the regex."""
     return [
         Range(
             start=Position(line=line_number, character=match.start()),
-            end=Position(line=line_number, character=match.end())
+            end=Position(line=line_number, character=match.end()),
         )
         for line_number, line in enumerate(read_file)
         for match in [m for m in [re.search(regex, line)] if m]
     ]
-
-