fix test cases

apache · Jan 18, 2024 · e3a4a20 · e3a4a20
1 parent b1a6a48
commit e3a4a20
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/tests/providers/snowflake/operators/test_snowflake.py b/tests/providers/snowflake/operators/test_snowflake.py
@@ -253,7 +253,9 @@ def test_snowflake_sql_api_to_fails_when_one_query_fails(
 
     @pytest.mark.parametrize("mock_sql, statement_count", [(SQL_MULTIPLE_STMTS, 4), (SINGLE_STMT, 1)])
     @mock.patch("airflow.providers.snowflake.hooks.snowflake_sql_api.SnowflakeSqlApiHook.execute_query")
-    def test_snowflake_sql_api_execute_operator_async(self, mock_db_hook, mock_sql, statement_count):
+    def test_snowflake_sql_api_execute_operator_async(
+        self, mock_execute_query, mock_sql, statement_count, mock_get_sql_api_query_status
+    ):
         """
         Asserts that a task is deferred and an SnowflakeSqlApiTrigger will be fired
         when the SnowflakeSqlApiOperator is executed.
@@ -266,6 +268,9 @@ def test_snowflake_sql_api_execute_operator_async(self, mock_db_hook, mock_sql,
             deferrable=True,
         )
 
+        mock_execute_query.return_value = ["uuid1"]
+        mock_get_sql_api_query_status.side_effect = [{"status": "running"}]
+
         with pytest.raises(TaskDeferred) as exc:
             operator.execute(create_context(operator))