subbota19 commented on code in PR #46672:
URL: https://github.com/apache/airflow/pull/46672#discussion_r1955909908
##########
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:
Hi @rawwar!
Updated. I was wondering if it's possible to use a pre-commit hook to handle
such docstring requirements. It could be beneficial to add something like this.
##########
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`
Review Comment:
Updated.
--
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]
