RDS Instance System Tests

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]

tests/system/providers/amazon/aws/rds/example_rds_instance.py → 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,56 +25,78 @@
     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"
+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.
 RDS_PASSWORD = 'database_password'
 
 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_class="db.t4g.micro",
-        engine="postgres",
+        db_instance_identifier=rds_db_identifier,
+        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",
-        db_identifier=RDS_DB_IDENTIFIER,
+    await_db_instance = RdsDbSensor(
+        task_id='await_db_instance',
+        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,
+            '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,
+        await_db_instance,
+        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