From 828d946d9b4f4db1d2fe831c5fc6c5d3ea579bea Mon Sep 17 00:00:00 2001
From: ferruzzi <ferruzzi@amazon.com>
Date: Wed, 21 Sep 2022 14:10:40 -0700
Subject: [PATCH 1/3] RDS Instance System Tests

---
 .../aws/{rds => }/example_rds_instance.py     | 36 ++++++++++++++-----
 .../providers/amazon/aws/rds/__init__.py      | 17 ---------
 2 files changed, 27 insertions(+), 26 deletions(-)
 rename tests/system/providers/amazon/aws/{rds => }/example_rds_instance.py (73%)
 delete mode 100644 tests/system/providers/amazon/aws/rds/__init__.py

diff --git a/tests/system/providers/amazon/aws/rds/example_rds_instance.py b/tests/system/providers/amazon/aws/example_rds_instance.py
similarity index 73%
rename from tests/system/providers/amazon/aws/rds/example_rds_instance.py
rename to tests/system/providers/amazon/aws/example_rds_instance.py
index ec6cba2d2550c..e364a3d16b270 100644
--- a/tests/system/providers/amazon/aws/rds/example_rds_instance.py
+++ b/tests/system/providers/amazon/aws/example_rds_instance.py
@@ -1,4 +1,3 @@
-#
 # Licensed to the Apache Software Foundation (ASF) under one
 # or more contributor license agreements.  See the NOTICE file
 # distributed with this work for additional information
@@ -26,12 +25,13 @@
     RdsDeleteDbInstanceOperator,
 )
 from airflow.providers.amazon.aws.sensors.rds import RdsDbSensor
-from tests.system.providers.amazon.aws.utils import set_env_id
+from airflow.utils.trigger_rule import TriggerRule
+from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder
+
+sys_test_context_task = SystemTestContextBuilder().build()
 
-ENV_ID = set_env_id()
 DAG_ID = "example_rds_instance"
 
-RDS_DB_IDENTIFIER = f'{ENV_ID}-database'
 RDS_USERNAME = 'database_username'
 # NEVER store your production password in plaintext in a DAG like this.
 # Use Airflow Secrets or a secret manager for this in production.
@@ -39,15 +39,18 @@
 
 with DAG(
     dag_id=DAG_ID,
-    schedule=None,
+    schedule='@once',
     start_date=datetime(2021, 1, 1),
     tags=['example'],
     catchup=False,
 ) as dag:
+    test_context = sys_test_context_task()
+    rds_db_identifier = f'{test_context[ENV_ID_KEY]}-database'
+
     # [START howto_operator_rds_create_db_instance]
     create_db_instance = RdsCreateDbInstanceOperator(
         task_id='create_db_instance',
-        db_instance_identifier=RDS_DB_IDENTIFIER,
+        db_instance_identifier=rds_db_identifier,
         db_instance_class="db.t4g.micro",
         engine="postgres",
         rds_kwargs={
@@ -61,21 +64,36 @@
     # [START howto_sensor_rds_instance]
     db_instance_available = RdsDbSensor(
         task_id="db_instance_available",
-        db_identifier=RDS_DB_IDENTIFIER,
+        db_identifier=rds_db_identifier,
     )
     # [END howto_sensor_rds_instance]
 
     # [START howto_operator_rds_delete_db_instance]
     delete_db_instance = RdsDeleteDbInstanceOperator(
         task_id='delete_db_instance',
-        db_instance_identifier=RDS_DB_IDENTIFIER,
+        db_instance_identifier=rds_db_identifier,
         rds_kwargs={
             "SkipFinalSnapshot": True,
         },
     )
     # [END howto_operator_rds_delete_db_instance]
+    delete_db_instance.trigger_rule = TriggerRule.ALL_DONE
+
+    chain(
+        # TEST SETUP
+        test_context,
+        # TEST BODY
+        create_db_instance,
+        db_instance_available,
+        delete_db_instance,
+    )
+
+    from tests.system.utils.watcher import watcher
+
+    # This test needs watcher in order to properly mark success/failure
+    # when "tearDown" task with trigger rule is part of the DAG
+    list(dag.tasks) >> watcher()
 
-    chain(create_db_instance, db_instance_available, delete_db_instance)
 
 from tests.system.utils import get_test_run  # noqa: E402
 
diff --git a/tests/system/providers/amazon/aws/rds/__init__.py b/tests/system/providers/amazon/aws/rds/__init__.py
deleted file mode 100644
index 217e5db960782..0000000000000
--- a/tests/system/providers/amazon/aws/rds/__init__.py
+++ /dev/null
@@ -1,17 +0,0 @@
-#
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements.  See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership.  The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License.  You may obtain a copy of the License at
-#
-#   http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied.  See the License for the
-# specific language governing permissions and limitations
-# under the License.

From 7ee6b6df4cba096f63e4e1fd1ce52f9e29b4eede Mon Sep 17 00:00:00 2001
From: ferruzzi <ferruzzi@amazon.com>
Date: Mon, 26 Sep 2022 13:44:56 -0700
Subject: [PATCH 2/3] RDS doc fixes

---
 docs/apache-airflow-providers-amazon/operators/rds.rst | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/docs/apache-airflow-providers-amazon/operators/rds.rst b/docs/apache-airflow-providers-amazon/operators/rds.rst
index e5f46ad66be42..d0752a014d0d3 100644
--- a/docs/apache-airflow-providers-amazon/operators/rds.rst
+++ b/docs/apache-airflow-providers-amazon/operators/rds.rst
@@ -146,7 +146,7 @@ Create a database instance
 To create a AWS DB instance you can use
 :class:`~airflow.providers.amazon.aws.operators.rds.RdsCreateDbInstanceOperator`.
 
-.. exampleinclude:: /../../tests/system/providers/amazon/aws/rds/example_rds_instance.py
+.. exampleinclude:: /../../tests/system/providers/amazon/aws/example_rds_instance.py
     :language: python
     :dedent: 4
     :start-after: [START howto_operator_rds_create_db_instance]
@@ -160,7 +160,7 @@ Delete a database instance
 To delete a AWS DB instance you can use
 :class:`~airflow.providers.amazon.aws.operators.rds.RDSDeleteDbInstanceOperator`.
 
-.. exampleinclude:: /../../tests/system/providers/amazon/aws/rds/example_rds_instance.py
+.. exampleinclude:: /../../tests/system/providers/amazon/aws/example_rds_instance.py
     :language: python
     :dedent: 4
     :start-after: [START howto_operator_rds_delete_db_instance]
@@ -178,7 +178,7 @@ To wait for an Amazon RDS instance or cluster to reach a specific status you can
 :class:`~airflow.providers.amazon.aws.sensors.rds.RdsDbSensor`.
 By default, the sensor waits for a database instance to reach the ``available`` state.
 
-.. exampleinclude:: /../../tests/system/providers/amazon/aws/rds/example_rds_instance.py
+.. exampleinclude:: /../../tests/system/providers/amazon/aws/example_rds_instance.py
     :language: python
     :dedent: 4
     :start-after: [START howto_sensor_rds_instance]

From d39d9e61bc3f7e8ca7bd31ee8ca33585dd911acf Mon Sep 17 00:00:00 2001
From: ferruzzi <ferruzzi@amazon.com>
Date: Mon, 26 Sep 2022 16:36:12 -0700
Subject: [PATCH 3/3] Set 'wait_for_completion = False' otherwise the Sensor is
 not being tested.

Also made some stylistic changes for consistency
---
 .../amazon/aws/example_rds_instance.py        | 23 +++++++++++--------
 1 file changed, 13 insertions(+), 10 deletions(-)

diff --git a/tests/system/providers/amazon/aws/example_rds_instance.py b/tests/system/providers/amazon/aws/example_rds_instance.py
index e364a3d16b270..c27f7f99eeedb 100644
--- a/tests/system/providers/amazon/aws/example_rds_instance.py
+++ b/tests/system/providers/amazon/aws/example_rds_instance.py
@@ -30,7 +30,7 @@
 
 sys_test_context_task = SystemTestContextBuilder().build()
 
-DAG_ID = "example_rds_instance"
+DAG_ID = 'example_rds_instance'
 
 RDS_USERNAME = 'database_username'
 # NEVER store your production password in plaintext in a DAG like this.
@@ -51,19 +51,22 @@
     create_db_instance = RdsCreateDbInstanceOperator(
         task_id='create_db_instance',
         db_instance_identifier=rds_db_identifier,
-        db_instance_class="db.t4g.micro",
-        engine="postgres",
+        db_instance_class='db.t4g.micro',
+        engine='postgres',
         rds_kwargs={
-            "MasterUsername": RDS_USERNAME,
-            "MasterUserPassword": RDS_PASSWORD,
-            "AllocatedStorage": 20,
+            'MasterUsername': RDS_USERNAME,
+            'MasterUserPassword': RDS_PASSWORD,
+            'AllocatedStorage': 20,
         },
     )
     # [END howto_operator_rds_create_db_instance]
 
+    # RdsCreateDbInstanceOperator waits by default, setting as False to test the Sensor below.
+    create_db_instance.wait_for_completion = False
+
     # [START howto_sensor_rds_instance]
-    db_instance_available = RdsDbSensor(
-        task_id="db_instance_available",
+    await_db_instance = RdsDbSensor(
+        task_id='await_db_instance',
         db_identifier=rds_db_identifier,
     )
     # [END howto_sensor_rds_instance]
@@ -73,7 +76,7 @@
         task_id='delete_db_instance',
         db_instance_identifier=rds_db_identifier,
         rds_kwargs={
-            "SkipFinalSnapshot": True,
+            'SkipFinalSnapshot': True,
         },
     )
     # [END howto_operator_rds_delete_db_instance]
@@ -84,7 +87,7 @@
         test_context,
         # TEST BODY
         create_db_instance,
-        db_instance_available,
+        await_db_instance,
         delete_db_instance,
     )
 
