@@ -177,19 +177,22 @@ def test_cloud_fetch(self):
177177for i in range (len (cf_result )):
178178assert cf_result [i ]== noop_result [i ]
179179
180- def test_execute_async (self ):
181- def isExecuting (operation_state ):
182- return not operation_state or operation_state in [
183- ttypes .TOperationState .RUNNING_STATE ,
184- ttypes .TOperationState .PENDING_STATE ,
185- ]
180+
181+ class TestPySQLAsyncQueriesSuite (PySQLPytestTestCase ):
182+ def isExecuting (self ,operation_state ):
183+ return not operation_state or operation_state in [
184+ ttypes .TOperationState .RUNNING_STATE ,
185+ ttypes .TOperationState .PENDING_STATE ,
186+ ]
187+
188+ def test_execute_async__long_running (self ):
186189
187190long_running_query = "SELECT COUNT(*) FROM RANGE(10000 * 16) x JOIN RANGE(10000) y ON FROM_UNIXTIME(x.id * y.id, 'yyyy-MM-dd') LIKE '%not%a%date%'"
188191with self .cursor ()as cursor :
189192cursor .execute_async (long_running_query )
190193
191194## Polling after every POLLING_INTERVAL seconds
192- while isExecuting (cursor .get_query_state ()):
195+ while self . isExecuting (cursor .get_query_state ()):
193196time .sleep (self .POLLING_INTERVAL )
194197log .info ("Polling the status in test_execute_async" )
195198
@@ -198,6 +201,55 @@ def isExecuting(operation_state):
198201
199202assert result [0 ].asDict ()== {"count(1)" :0 }
200203
204+ def test_execute_async__small_result (self ):
205+ small_result_query = "SELECT 1"
206+
207+ with self .cursor ()as cursor :
208+ cursor .execute_async (small_result_query )
209+
210+ ## Fake sleep for 5 secs
211+ time .sleep (5 )
212+
213+ ## Polling after every POLLING_INTERVAL seconds
214+ while self .isExecuting (cursor .get_query_state ()):
215+ time .sleep (self .POLLING_INTERVAL )
216+ log .info ("Polling the status in test_execute_async" )
217+
218+ cursor .get_async_execution_result ()
219+ result = cursor .fetchall ()
220+
221+ assert result [0 ].asDict ()== {"1" :1 }
222+
223+ def test_execute_async__large_result (self ):
224+ x_dimension = 1000
225+ y_dimension = 1000
226+ large_result_query = f"""
227+ SELECT
228+ x.id AS x_id,
229+ y.id AS y_id,
230+ FROM_UNIXTIME(x.id * y.id, 'yyyy-MM-dd') AS date
231+ FROM
232+ RANGE({ x_dimension }
233+ JOIN
234+ RANGE({ y_dimension }
235+ """
236+
237+ with self .cursor ()as cursor :
238+ cursor .execute_async (large_result_query )
239+
240+ ## Fake sleep for 5 secs
241+ time .sleep (5 )
242+
243+ ## Polling after every POLLING_INTERVAL seconds
244+ while self .isExecuting (cursor .get_query_state ()):
245+ time .sleep (self .POLLING_INTERVAL )
246+ log .info ("Polling the status in test_execute_async" )
247+
248+ cursor .get_async_execution_result ()
249+ result = cursor .fetchall ()
250+
251+ assert len (result )== x_dimension * y_dimension
252+
201253
202254# Exclude Retry tests because they require specific setups, and LargeQueries too slow for core
203255# tests