rawwar commented on code in PR #46672:
URL: https://github.com/apache/airflow/pull/46672#discussion_r1955604794
##########
providers/snowflake/tests/provider_tests/snowflake/operators/test_snowflake.py:
##########
@@ -391,3 +391,66 @@ def
test_snowflake_sql_api_execute_operator_running_before_defer(
operator.execute(create_context(operator))
assert mock_defer.called
+
+ def test_snowflake_sql_api_execute_operator_polling_running(
+ self, mock_execute_query, mock_get_sql_api_query_status,
mock_check_query_output
+ ):
+ """
+ Tests that the execute method correctly loops and waits until all
queries complete
+ when `deferrable=False`
+ """
+ operator = SnowflakeSqlApiOperator(
+ task_id=TASK_ID,
+ snowflake_conn_id=CONN_ID,
+ sql=SQL_MULTIPLE_STMTS,
+ statement_count=4,
+ do_xcom_push=False,
+ deferrable=False,
+ )
+
+ mock_execute_query.return_value = ["uuid1"]
+
+ mock_get_sql_api_query_status.side_effect = [
+ # Initial get_sql_api_query_status check
+ {"status": "running"},
+ # 1st poll_on_queries check (poll_interval: 5s)
+ {"status": "running"},
+ # 2nd poll_on_queries check (poll_interval: 5s)
+ {"status": "running"},
+ # 3rd poll_on_queries check (poll_interval: 5s)
+ {"status": "success"},
+ ]
+
+ with mock.patch("time.sleep") as mock_sleep:
+ operator.execute(context=None)
+ mock_check_query_output.assert_called_once_with(["uuid1"])
+ assert mock_sleep.call_count == 3
+
+ def test_snowflake_sql_api_execute_operator_polling_failed(
+ self, mock_execute_query, mock_get_sql_api_query_status,
mock_check_query_output
+ ):
+ """
+ Tests that the execute method raises AirflowException if any query
fails during polling
+ when `deferrable=False`
Review Comment:
```suggestion
when ``deferrable=False``
```
we need to use 2 backticks on either sides for inline literals. refer to any
other method's docstring
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
