From fcf1e923598b125baf16c229309959893683e11b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rapha=C3=ABl=20Vandon?= <vandonr@amazon.com>
Date: Mon, 30 Jan 2023 18:02:54 -0800
Subject: [PATCH] shorten other wait times in sys tests

---
 airflow/providers/amazon/aws/operators/rds.py               | 6 +++++-
 tests/system/providers/amazon/aws/example_batch.py          | 2 ++
 .../amazon/aws/example_eks_with_fargate_in_one_step.py      | 1 +
 .../amazon/aws/example_eks_with_fargate_profile.py          | 2 ++
 .../amazon/aws/example_eks_with_nodegroup_in_one_step.py    | 2 ++
 .../providers/amazon/aws/example_eks_with_nodegroups.py     | 2 ++
 tests/system/providers/amazon/aws/example_emr.py            | 1 +
 tests/system/providers/amazon/aws/example_emr_eks.py        | 2 ++
 tests/system/providers/amazon/aws/example_emr_serverless.py | 3 +++
 tests/system/providers/amazon/aws/example_glue.py           | 1 +
 tests/system/providers/amazon/aws/example_rds_export.py     | 1 +
 11 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/airflow/providers/amazon/aws/operators/rds.py b/airflow/providers/amazon/aws/operators/rds.py
index 2f2cf58438c22..a3d253203242a 100644
--- a/airflow/providers/amazon/aws/operators/rds.py
+++ b/airflow/providers/amazon/aws/operators/rds.py
@@ -370,6 +370,7 @@ def __init__(
         *,
         export_task_identifier: str,
         wait_for_completion: bool = True,
+        check_interval: int = 30,
         aws_conn_id: str = "aws_default",
         **kwargs,
     ):
@@ -377,6 +378,7 @@ def __init__(
 
         self.export_task_identifier = export_task_identifier
         self.wait_for_completion = wait_for_completion
+        self.check_interval = check_interval
 
     def execute(self, context: Context) -> str:
         self.log.info("Canceling export task %s", self.export_task_identifier)
@@ -386,7 +388,9 @@ def execute(self, context: Context) -> str:
         )
 
         if self.wait_for_completion:
-            self.hook.wait_for_export_task_state(self.export_task_identifier, target_state="canceled")
+            self.hook.wait_for_export_task_state(
+                self.export_task_identifier, target_state="canceled", check_interval=self.check_interval
+            )
         return json.dumps(cancel_export, default=str)
 
 
diff --git a/tests/system/providers/amazon/aws/example_batch.py b/tests/system/providers/amazon/aws/example_batch.py
index 1937127fcd33c..a035b12f87b6e 100644
--- a/tests/system/providers/amazon/aws/example_batch.py
+++ b/tests/system/providers/amazon/aws/example_batch.py
@@ -207,6 +207,7 @@ def delete_job_queue(job_queue_name):
         job_id=submit_batch_job.output,
     )
     # [END howto_sensor_batch]
+    wait_for_batch_job.poke_interval = 10
 
     wait_for_compute_environment_disabled = BatchComputeEnvironmentSensor(
         task_id="wait_for_compute_environment_disabled",
@@ -224,6 +225,7 @@ def delete_job_queue(job_queue_name):
         task_id="wait_for_job_queue_deleted",
         job_queue=batch_job_queue_name,
         treat_non_existing_as_deleted=True,
+        poke_interval=10,
     )
 
     log_cleanup = prune_logs(
diff --git a/tests/system/providers/amazon/aws/example_eks_with_fargate_in_one_step.py b/tests/system/providers/amazon/aws/example_eks_with_fargate_in_one_step.py
index 8c96ba15cc85b..01613886de97c 100644
--- a/tests/system/providers/amazon/aws/example_eks_with_fargate_in_one_step.py
+++ b/tests/system/providers/amazon/aws/example_eks_with_fargate_in_one_step.py
@@ -137,6 +137,7 @@
         trigger_rule=TriggerRule.ALL_DONE,
         cluster_name=cluster_name,
         target_state=ClusterStates.NONEXISTENT,
+        poke_interval=10,
     )
 
     chain(
diff --git a/tests/system/providers/amazon/aws/example_eks_with_fargate_profile.py b/tests/system/providers/amazon/aws/example_eks_with_fargate_profile.py
index 72e273645cb94..d3862f9d3751e 100644
--- a/tests/system/providers/amazon/aws/example_eks_with_fargate_profile.py
+++ b/tests/system/providers/amazon/aws/example_eks_with_fargate_profile.py
@@ -157,6 +157,7 @@
         fargate_profile_name=fargate_profile_name,
         target_state=FargateProfileStates.NONEXISTENT,
         trigger_rule=TriggerRule.ALL_DONE,
+        poke_interval=10,
     )
 
     delete_cluster = EksDeleteClusterOperator(
@@ -169,6 +170,7 @@
         task_id="wait_for_delete_cluster",
         cluster_name=cluster_name,
         target_state=ClusterStates.NONEXISTENT,
+        poke_interval=10,
     )
 
     chain(
diff --git a/tests/system/providers/amazon/aws/example_eks_with_nodegroup_in_one_step.py b/tests/system/providers/amazon/aws/example_eks_with_nodegroup_in_one_step.py
index f90288a65ded5..e565eb9f78cc8 100644
--- a/tests/system/providers/amazon/aws/example_eks_with_nodegroup_in_one_step.py
+++ b/tests/system/providers/amazon/aws/example_eks_with_nodegroup_in_one_step.py
@@ -75,6 +75,7 @@
         cluster_name=cluster_name,
         nodegroup_name=nodegroup_name,
         target_state=NodegroupStates.ACTIVE,
+        poke_interval=10,
     )
 
     start_pod = EksPodOperator(
@@ -105,6 +106,7 @@
         trigger_rule=TriggerRule.ALL_DONE,
         cluster_name=cluster_name,
         target_state=ClusterStates.NONEXISTENT,
+        poke_interval=10,
     )
 
     chain(
diff --git a/tests/system/providers/amazon/aws/example_eks_with_nodegroups.py b/tests/system/providers/amazon/aws/example_eks_with_nodegroups.py
index 07f61ae38a9f4..bb9f9c21c95e3 100644
--- a/tests/system/providers/amazon/aws/example_eks_with_nodegroups.py
+++ b/tests/system/providers/amazon/aws/example_eks_with_nodegroups.py
@@ -95,6 +95,7 @@
         target_state=NodegroupStates.ACTIVE,
     )
     # [END howto_sensor_eks_nodegroup]
+    await_create_nodegroup.poke_interval = 10
 
     # [START howto_operator_eks_pod_operator]
     start_pod = EksPodOperator(
@@ -140,6 +141,7 @@
         trigger_rule=TriggerRule.ALL_DONE,
         cluster_name=cluster_name,
         target_state=ClusterStates.NONEXISTENT,
+        poke_interval=10,
     )
 
     chain(
diff --git a/tests/system/providers/amazon/aws/example_emr.py b/tests/system/providers/amazon/aws/example_emr.py
index f1d8665f3298a..dbc1430f01f44 100644
--- a/tests/system/providers/amazon/aws/example_emr.py
+++ b/tests/system/providers/amazon/aws/example_emr.py
@@ -156,6 +156,7 @@ def delete_security_config(config_name: str):
     # [START howto_sensor_emr_job_flow]
     check_job_flow = EmrJobFlowSensor(task_id="check_job_flow", job_flow_id=create_job_flow.output)
     # [END howto_sensor_emr_job_flow]
+    check_job_flow.poke_interval = 10
 
     delete_security_configuration = delete_security_config(config_name)
 
diff --git a/tests/system/providers/amazon/aws/example_emr_eks.py b/tests/system/providers/amazon/aws/example_emr_eks.py
index d7465a79e6d10..16d1f3bc2475a 100644
--- a/tests/system/providers/amazon/aws/example_emr_eks.py
+++ b/tests/system/providers/amazon/aws/example_emr_eks.py
@@ -210,6 +210,7 @@ def delete_virtual_cluster(virtual_cluster_id):
         cluster_name=eks_cluster_name,
         nodegroup_name=nodegroup_name,
         target_state=NodegroupStates.ACTIVE,
+        poke_interval=10,
     )
 
     emr_access_on_eks = enable_access_emr_on_eks(eks_cluster_name, eks_namespace)
@@ -260,6 +261,7 @@ def delete_virtual_cluster(virtual_cluster_id):
         cluster_name=eks_cluster_name,
         target_state=ClusterStates.NONEXISTENT,
         trigger_rule=TriggerRule.ALL_DONE,
+        poke_interval=10,
     )
 
     delete_bucket = S3DeleteBucketOperator(
diff --git a/tests/system/providers/amazon/aws/example_emr_serverless.py b/tests/system/providers/amazon/aws/example_emr_serverless.py
index c2838f85859d5..6d8a669c3e97e 100644
--- a/tests/system/providers/amazon/aws/example_emr_serverless.py
+++ b/tests/system/providers/amazon/aws/example_emr_serverless.py
@@ -88,6 +88,7 @@
         application_id=emr_serverless_app_id,
     )
     # [END howto_sensor_emr_serverless_application]
+    wait_for_app_creation.poke_interval = 1
 
     # [START howto_operator_emr_serverless_start_job]
     start_job = EmrServerlessStartJobOperator(
@@ -98,6 +99,7 @@
         configuration_overrides=SPARK_CONFIGURATION_OVERRIDES,
     )
     # [END howto_operator_emr_serverless_start_job]
+    start_job.waiter_check_interval_seconds = 10
 
     # [START howto_sensor_emr_serverless_job]
     wait_for_job = EmrServerlessJobSensor(
@@ -113,6 +115,7 @@
         application_id=emr_serverless_app_id,
     )
     # [END howto_operator_emr_serverless_delete_application]
+    delete_app.waiter_check_interval_seconds = 1
     delete_app.trigger_rule = TriggerRule.ALL_DONE
 
     delete_s3_bucket = S3DeleteBucketOperator(
diff --git a/tests/system/providers/amazon/aws/example_glue.py b/tests/system/providers/amazon/aws/example_glue.py
index 15edcfacb044f..f010b2dfd9fe0 100644
--- a/tests/system/providers/amazon/aws/example_glue.py
+++ b/tests/system/providers/amazon/aws/example_glue.py
@@ -164,6 +164,7 @@ def glue_cleanup(crawler_name: str, job_name: str, db_name: str) -> None:
         run_id=submit_glue_job.output,
     )
     # [END howto_sensor_glue]
+    wait_for_job.poke_interval = 10
 
     delete_bucket = S3DeleteBucketOperator(
         task_id="delete_bucket",
diff --git a/tests/system/providers/amazon/aws/example_rds_export.py b/tests/system/providers/amazon/aws/example_rds_export.py
index 2bdf02b906e9f..d103763952417 100644
--- a/tests/system/providers/amazon/aws/example_rds_export.py
+++ b/tests/system/providers/amazon/aws/example_rds_export.py
@@ -127,6 +127,7 @@ def get_snapshot_arn(snapshot_name: str) -> str:
         export_task_identifier=rds_export_task_id,
     )
     # [END howto_operator_rds_cancel_export]
+    cancel_export.check_interval = 10
 
     # [START howto_sensor_rds_export_task_existence]
     export_sensor = RdsExportTaskExistenceSensor(
