1+ from lsprotocol .types import Position
2+ from sqlmesh .core .context import Context
3+ from sqlmesh .lsp .context import LSPContext , ModelTarget
4+ from sqlmesh .lsp .reference import (
5+ get_macro_find_all_references ,
6+ get_macro_definitions_for_a_path ,
7+ )
8+ from sqlmesh .lsp .uri import URI
9+
10+
11+ def test_find_all_references_for_macro_add_one ():
12+ """Test finding all references to the @ADD_ONE macro."""
13+ context = Context (paths = ["examples/sushi" ])
14+ lsp_context = LSPContext (context )
15+
16+ # Find the top_waiters model that uses @ADD_ONE macro
17+ top_waiters_path = next (
18+ path
19+ for path , info in lsp_context .map .items ()
20+ if isinstance (info , ModelTarget ) and "sushi.top_waiters" in info .names
21+ )
22+
23+ top_waiters_uri = URI .from_path (top_waiters_path )
24+ macro_references = get_macro_definitions_for_a_path (lsp_context , top_waiters_uri )
25+
26+ # Find the @ADD_ONE reference
27+ add_one_ref = next ((ref for ref in macro_references if ref .range .start .line == 12 ), None )
28+ assert add_one_ref is not None , "Should find @ADD_ONE reference in top_waiters"
29+
30+ # Click on the @ADD_ONE macro at line 13, character 5 (the @ symbol)
31+ position = Position (line = 12 , character = 5 )
32+
33+ all_references = get_macro_find_all_references (lsp_context , top_waiters_uri , position )
34+
35+ # Should find at least 2 references: the definition and the usage in top_waiters
36+ assert len (all_references ) >= 2 , f"Expected at least 2 references, found { len (all_references )}
37+
38+ # Verify the macro definition is included
39+ definition_refs = [ref for ref in all_references if "utils.py" in ref .uri ]
40+ assert len (definition_refs ) >= 1 , "Should include the macro definition in utils.py"
41+
42+ # Verify the usage in top_waiters is included
43+ usage_refs = [ref for ref in all_references if "top_waiters" in ref .uri ]
44+ assert len (usage_refs ) >= 1 , "Should include the usage in top_waiters.sql"
45+
46+ # breakpoint()
47+ expected_ranges = [
48+ # Macro definition in utils.py
49+ {
50+ "uri" : "file:///Users/themistoklisvaltinos/Developer/sqlmesh/examples/sushi/macros/utils.py" ,
51+ "range" : ((6 , 0 ), (9 , 22 )),
52+ },
53+ # Usage in customers.sql
54+ {
55+ "uri" : "file:///Users/themistoklisvaltinos/Developer/sqlmesh/examples/sushi/models/customers.sql" ,
56+ "range" : ((36 , 7 ), (36 , 14 )),
57+ },
58+ # Usage in top_waiters.sql
59+ {
60+ "uri" : "file:///Users/themistoklisvaltinos/Developer/sqlmesh/examples/sushi/models/top_waiters.sql" ,
61+ "range" : ((12 , 5 ), (12 , 12 )),
62+ },
63+ ]
64+
65+ for expected in expected_ranges :
66+ assert any (
67+ ref .uri == expected ["uri" ] and
68+ ref .range .start .line == expected ["range" ][0 ][0 ] and
69+ ref .range .start .character == expected ["range" ][0 ][1 ] and
70+ ref .range .end .line == expected ["range" ][1 ][0 ] and
71+ ref .range .end .character == expected ["range" ][1 ][1 ]
72+ for ref in all_references
73+ ), f"Expected reference with uri { expected ['uri' ]} { expected ['range' ]}
74+
75+
76+ def test_find_all_references_for_macro_multiply ():
77+ """Test finding all references to the @MULTIPLY macro."""
78+ context = Context (paths = ["examples/sushi" ])
79+ lsp_context = LSPContext (context )
80+
81+ # Find the top_waiters model that uses @MULTIPLY macro
82+ top_waiters_path = next (
83+ path
84+ for path , info in lsp_context .map .items ()
85+ if isinstance (info , ModelTarget ) and "sushi.top_waiters" in info .names
86+ )
87+
88+ top_waiters_uri = URI .from_path (top_waiters_path )
89+ macro_references = get_macro_definitions_for_a_path (lsp_context , top_waiters_uri )
90+
91+ # Find the @MULTIPLY reference
92+ multiply_ref = next ((ref for ref in macro_references if ref .range .start .line == 13 ), None )
93+ assert multiply_ref is not None , "Should find @MULTIPLY reference in top_waiters"
94+
95+ # Click on the @MULTIPLY macro at line 14, character 5 (the @ symbol)
96+ position = Position (line = 13 , character = 5 )
97+ all_references = get_macro_find_all_references (lsp_context , top_waiters_uri , position )
98+
99+ # Should find at least 2 references: the definition and the usage
100+ assert len (all_references ) >= 2 , f"Expected at least 2 references, found { len (all_references )}
101+
102+ # Verify both definition and usage are included
103+ assert any ("utils.py" in ref .uri for ref in all_references ), "Should include macro definition"
104+ assert any ("top_waiters" in ref .uri for ref in all_references ), "Should include usage"
105+
106+
107+ def test_find_all_references_for_sql_literal_macro ():
108+ """Test finding references to @SQL_LITERAL macro ."""
109+ context = Context (paths = ["examples/sushi" ])
110+ lsp_context = LSPContext (context )
111+
112+ # Find the top_waiters model that uses @SQL_LITERAL macro
113+ top_waiters_path = next (
114+ path
115+ for path , info in lsp_context .map .items ()
116+ if isinstance (info , ModelTarget ) and "sushi.top_waiters" in info .names
117+ )
118+
119+ top_waiters_uri = URI .from_path (top_waiters_path )
120+ macro_references = get_macro_definitions_for_a_path (lsp_context , top_waiters_uri )
121+
122+ # Find the @SQL_LITERAL reference
123+ sql_literal_ref = next ((ref for ref in macro_references if ref .range .start .line == 14 ), None )
124+ assert sql_literal_ref is not None , "Should find @SQL_LITERAL reference in top_waiters"
125+
126+ # Click on the @SQL_LITERAL macro
127+ position = Position (line = 14 , character = 5 )
128+ all_references = get_macro_find_all_references (lsp_context , top_waiters_uri , position )
129+
130+ # For user-defined macros in utils.py, should find references
131+ assert len (all_references ) >= 2 , f"Expected at least 2 references, found { len (all_references )}
132+
133+
134+ def test_find_references_from_outside_macro_position ():
135+ """Test that clicking outside a macro doesn't return macro references."""
136+ context = Context (paths = ["examples/sushi" ])
137+ lsp_context = LSPContext (context )
138+
139+ top_waiters_path = next (
140+ path
141+ for path , info in lsp_context .map .items ()
142+ if isinstance (info , ModelTarget ) and "sushi.top_waiters" in info .names
143+ )
144+
145+ top_waiters_uri = URI .from_path (top_waiters_path )
146+
147+ # Click on a position that is not on a macro
148+ position = Position (line = 0 , character = 0 ) # First line, which is a comment
149+ all_references = get_macro_find_all_references (lsp_context , top_waiters_uri , position )
150+
151+ # Should return empty list when not on a macro
152+ assert len (all_references ) == 0 , "Should not find macro references when not on a macro"
153+
154+
155+ def test_multi_repo_macro_references ():
156+ """Test finding macro references across multiple repositories."""
157+ context = Context (paths = ["examples/multi/repo_1" , "examples/multi/repo_2" ], gateway = "memory" )
158+ lsp_context = LSPContext (context )
159+
160+ # Find model 'd' which uses macros from repo_2
161+ d_path = next (
162+ path
163+ for path , info in lsp_context .map .items ()
164+ if isinstance (info , ModelTarget ) and "silver.d" in info .names
165+ )
166+
167+ d_uri = URI .from_path (d_path )
168+ macro_references = get_macro_definitions_for_a_path (lsp_context , d_uri )
169+
170+ if macro_references :
171+ # Click on the second macro reference which appears under the same name in repo_1 ('dup')
172+ first_ref = macro_references [1 ]
173+ position = Position (
174+ line = first_ref .range .start .line ,
175+ character = first_ref .range .start .character + 1
176+ )
177+ all_references = get_macro_find_all_references (lsp_context , d_uri , position )
178+
179+ # Should find the definition and usage
180+ assert len (all_references ) == 2 , f"Expected 2 references, found { len (all_references )}
181+
182+ # Verify references from repo_2
183+ assert any ("repo_2" in ref .uri for ref in all_references ), "Should find macro in repo_2"
184+
185+ # But not references in repo_1 since despite identical name they're different macros
186+ assert not any ("repo_1" in ref .uri for ref in all_references ), "Shouldn't find macro in repo_1"
0 commit comments