@@ -53,6 +53,7 @@ def test_get_col_name(self) -> None:
5353 self .assertEqual (AGEGraph ._get_col_name (* value ), expected [idx ])
5454
5555 def test_wrap_query (self ) -> None :
56+ """Test basic query wrapping functionality."""
5657 inputs = [
5758 # Positive case: Simple return clause
5859 """
@@ -76,46 +77,195 @@ def test_wrap_query(self) -> None:
7677
7778 expected = [
7879 # Expected output for the first positive case
80+ """
81+ SELECT * FROM ag_catalog.cypher('test', $$
82+ MATCH (keanu:Person {name:'Keanu Reeves'})
83+ RETURN keanu.name AS name, keanu.born AS born
84+ $$) AS (name agtype, born agtype);
85+ """ ,
86+ # Second test case (no RETURN clause)
87+ """
88+ SELECT * FROM ag_catalog.cypher('test', $$
89+ MERGE (n:a {id: 1})
90+ $$) AS (a agtype);
91+ """ ,
92+ # Expected output for the negative cases (no RETURN clause)
93+ """
94+ SELECT * FROM ag_catalog.cypher('test', $$
95+ MATCH (n {description: "This will return a value"})
96+ MERGE (n)-[:RELATED]->(m)
97+ $$) AS (a agtype);
98+ """ ,
99+ """
100+ SELECT * FROM ag_catalog.cypher('test', $$
101+ MATCH (n {returnValue: "some value"})
102+ MERGE (n)-[:RELATED]->(m)
103+ $$) AS (a agtype);
104+ """ ,
105+ ]
106+
107+ for idx , value in enumerate (inputs ):
108+ result = AGEGraph ._wrap_query (value , "test" )
109+ expected_result = expected [idx ]
110+ self .assertEqual (
111+ re .sub (r"\s" , "" , result ),
112+ re .sub (r"\s" , "" , expected_result ),
113+ (
114+ f"Failed on test case { idx + 1 } \n "
115+ f"Input:\n { value } \n "
116+ f"Expected:\n { expected_result } \n "
117+ f"Got:\n { result }
118+ ),
119+ )
120+
121+ def test_wrap_query_union_except (self ) -> None :
122+ """Test query wrapping with UNION and EXCEPT operators."""
123+ inputs = [
124+ # UNION case
125+ """
126+ MATCH (n:Person)
127+ RETURN n.name AS name, n.age AS age
128+ UNION
129+ MATCH (n:Employee)
130+ RETURN n.name AS name, n.salary AS salary
131+ """ ,
132+ """
133+ MATCH (a:Employee {name: "Alice"})
134+ RETURN a.name AS name
135+ UNION
136+ MATCH (b:Manager {name: "Bob"})
137+ RETURN b.name AS name
138+ """ ,
139+ # Complex UNION case
140+ """
141+ MATCH (n)-[r]->(m)
142+ RETURN n.name AS source, type(r) AS relationship, m.name AS target
143+ UNION
144+ MATCH (m)-[r]->(n)
145+ RETURN m.name AS source, type(r) AS relationship, n.name AS target
146+ """ ,
147+ """
148+ MATCH (a:Person)-[:FRIEND]->(b:Person)
149+ WHERE a.age > 30
150+ RETURN a.name AS name
151+ UNION
152+ MATCH (c:Person)-[:FRIEND]->(d:Person)
153+ WHERE c.age < 25
154+ RETURN c.name AS name
155+ """ ,
156+ # EXCEPT case
157+ """
158+ MATCH (n:Person)
159+ RETURN n.name AS name
160+ EXCEPT
161+ MATCH (n:Employee)
162+ RETURN n.name AS name
163+ """ ,
164+ """
165+ MATCH (a:Person)
166+ RETURN a.name AS name, a.age AS age
167+ EXCEPT
168+ MATCH (b:Person {name: "Alice", age: 30})
169+ RETURN b.name AS name, b.age AS age
170+ """ ,
171+ ]
172+
173+ expected = [
79174 """
80175 SELECT * FROM ag_catalog.cypher('test', $$
81- MATCH (keanu:Person {name:'Keanu Reeves'})
82- RETURN keanu.name AS name, keanu.born AS born
83- $$) AS (name agtype, born agtype);
176+ MATCH (n:Person)
177+ RETURN n.name AS name, n.age AS age
178+ UNION
179+ MATCH (n:Employee)
180+ RETURN n.name AS name, n.salary AS salary
181+ $$) AS (name agtype, age agtype, salary agtype);
84182 """ ,
85183 """
86184 SELECT * FROM ag_catalog.cypher('test', $$
87- MERGE (n:a {id: 1})
88- $$) AS (a agtype);
185+ MATCH (a:Employee {name: "Alice"})
186+ RETURN a.name AS name
187+ UNION
188+ MATCH (b:Manager {name: "Bob"})
189+ RETURN b.name AS name
190+ $$) AS (name agtype);
89191 """ ,
90- # Expected output for the negative cases (no return clause)
91192 """
92193 SELECT * FROM ag_catalog.cypher('test', $$
93- MATCH (n {description: "This will return a value"})
94- MERGE (n)-[:RELATED]->(m)
95- $$) AS (a agtype);
194+ MATCH (n)-[r]->(m)
195+ RETURN n.name AS source, type(r) AS relationship, m.name AS target
196+ UNION
197+ MATCH (m)-[r]->(n)
198+ RETURN m.name AS source, type(r) AS relationship, n.name AS target
199+ $$) AS (source agtype, relationship agtype, target agtype);
96200 """ ,
97201 """
98202 SELECT * FROM ag_catalog.cypher('test', $$
99- MATCH (n {returnValue: "some value"})
100- MERGE (n)-[:RELATED]->(m)
101- $$) AS (a agtype);
203+ MATCH (a:Person)-[:FRIEND]->(b:Person)
204+ WHERE a.age > 30
205+ RETURN a.name AS name
206+ UNION
207+ MATCH (c:Person)-[:FRIEND]->(d:Person)
208+ WHERE c.age < 25
209+ RETURN c.name AS name
210+ $$) AS (name agtype);
211+ """ ,
212+ """
213+ SELECT * FROM ag_catalog.cypher('test', $$
214+ MATCH (n:Person)
215+ RETURN n.name AS name
216+ EXCEPT
217+ MATCH (n:Employee)
218+ RETURN n.name AS name
219+ $$) AS (name agtype);
220+ """ ,
221+ """
222+ SELECT * FROM ag_catalog.cypher('test', $$
223+ MATCH (a:Person)
224+ RETURN a.name AS name, a.age AS age
225+ EXCEPT
226+ MATCH (b:Person {name: "Alice", age: 30})
227+ RETURN b.name AS name, b.age AS age
228+ $$) AS (name agtype, age agtype);
102229 """ ,
103230 ]
104231
105232 for idx , value in enumerate (inputs ):
233+ result = AGEGraph ._wrap_query (value , "test" )
234+ expected_result = expected [idx ]
106235 self .assertEqual (
107- re .sub (r"\s" , "" , AGEGraph ._wrap_query (value , "test" )),
108- re .sub (r"\s" , "" , expected [idx ]),
236+ re .sub (r"\s" , "" , result ),
237+ re .sub (r"\s" , "" , expected_result ),
238+ (
239+ f"Failed on test case { idx + 1 } \n "
240+ f"Input:\n { value } \n "
241+ f"Expected:\n { expected_result } \n "
242+ f"Got:\n { result }
243+ ),
109244 )
110245
111- with self .assertRaises (ValueError ):
112- AGEGraph ._wrap_query (
113- """
246+ def test_wrap_query_errors (self ) -> None :
247+ """Test error cases for query wrapping."""
248+ error_cases = [
249+ # Empty query
250+ "" ,
251+ # Return * case
252+ """
114253 MATCH ()
115254 RETURN *
116255 """ ,
117- "test" ,
118- )
256+ # Return * in UNION
257+ """
258+ MATCH (n:Person)
259+ RETURN n.name
260+ UNION
261+ MATCH ()
262+ RETURN *
263+ """ ,
264+ ]
265+
266+ for query in error_cases :
267+ with self .assertRaises (ValueError ):
268+ AGEGraph ._wrap_query (query , "test" )
119269
120270 def test_format_properties (self ) -> None :
121271 inputs : List [Dict [str , Any ]] = [{}, {"a" : "b" }, {"a" : "b" , "c" : 1 , "d" : True }]
0 commit comments