Migrate Google example compute_igm to new design AIP-47 (#25132)

related: #22447, #22430
apache · Jul 19, 2022 · 9644451 · 9644451
1 parent f36f81e
commit 9644451
Show file tree

Hide file tree

Showing 5 changed files with 46 additions and 37 deletions.
diff --git a/TESTING.rst b/TESTING.rst
@@ -1100,8 +1100,7 @@ A simple example of a system test is available in:
 
 ``tests/providers/google/cloud/operators/test_compute_system.py``.
 
-It runs two DAGs defined in ``airflow.providers.google.cloud.example_dags.example_compute.py`` and
-``airflow.providers.google.cloud.example_dags.example_compute_igm.py``.
+It runs two DAGs defined in ``airflow.providers.google.cloud.example_dags.example_compute.py``.
 
 Preparing provider packages for System Tests for Airflow 1.10.* series
 ----------------------------------------------------------------------

diff --git a/docs/apache-airflow-providers-google/operators/cloud/compute.rst b/docs/apache-airflow-providers-google/operators/cloud/compute.rst
@@ -181,12 +181,12 @@ Using the operator
 
 The code to create the operator:
 
-.. exampleinclude:: /../../airflow/providers/google/cloud/example_dags/example_compute_igm.py
+.. exampleinclude:: /../../tests/system/providers/google/cloud/compute_igm/example_compute_igm.py
  :language: python
  :start-after: [START howto_operator_compute_template_copy_args]
  :end-before: [END howto_operator_compute_template_copy_args]
 
-.. exampleinclude:: /../../airflow/providers/google/cloud/example_dags/example_compute_igm.py
+.. exampleinclude:: /../../tests/system/providers/google/cloud/compute_igm/example_compute_igm.py
  :language: python
  :dedent: 4
  :start-after: [START howto_operator_gce_igm_copy_template]
@@ -195,7 +195,7 @@ The code to create the operator:
 You can also create the operator without project id - project id will be retrieved
 from the Google Cloud connection used:
 
-.. exampleinclude:: /../../airflow/providers/google/cloud/example_dags/example_compute_igm.py
+.. exampleinclude:: /../../tests/system/providers/google/cloud/compute_igm/example_compute_igm.py
  :language: python
  :dedent: 4
  :start-after: [START howto_operator_gce_igm_copy_template_no_project_id]
@@ -235,12 +235,12 @@ Using the operator
 
 The code to create the operator:
 
-.. exampleinclude:: /../../airflow/providers/google/cloud/example_dags/example_compute_igm.py
+.. exampleinclude:: /../../tests/system/providers/google/cloud/compute_igm/example_compute_igm.py
  :language: python
  :start-after: [START howto_operator_compute_igm_update_template_args]
  :end-before: [END howto_operator_compute_igm_update_template_args]
 
-.. exampleinclude:: /../../airflow/providers/google/cloud/example_dags/example_compute_igm.py
+.. exampleinclude:: /../../tests/system/providers/google/cloud/compute_igm/example_compute_igm.py
  :language: python
  :dedent: 4
  :start-after: [START howto_operator_gce_igm_update_template]
@@ -249,7 +249,7 @@ The code to create the operator:
 You can also create the operator without project id - project id will be retrieved
 from the Google Cloud connection used:
 
-.. exampleinclude:: /../../airflow/providers/google/cloud/example_dags/example_compute_igm.py
+.. exampleinclude:: /../../tests/system/providers/google/cloud/compute_igm/example_compute_igm.py
  :language: python
  :dedent: 4
  :start-after: [START howto_operator_gce_igm_update_template_no_project_id]

diff --git a/tests/providers/google/cloud/operators/test_compute_system.py b/tests/providers/google/cloud/operators/test_compute_system.py
@@ -41,24 +41,3 @@ def tearDown(self):
  @provide_gcp_context(GCP_COMPUTE_KEY)
  def test_run_example_dag_compute(self):
  self.run_dag('example_gcp_compute', CLOUD_DAG_FOLDER)
-
-
-@pytest.mark.backend("mysql", "postgres")
-@pytest.mark.credential_file(GCP_COMPUTE_KEY)
-class GcpComputeIgmExampleDagsSystemTest(GoogleSystemTest):
- helper = GCPComputeTestHelper()
-
- @provide_gcp_context(GCP_COMPUTE_KEY)
- def setUp(self):
- super().setUp()
- self.helper.delete_instance_group_and_template(silent=True)
- self.helper.create_instance_group_and_template()
-
- @provide_gcp_context(GCP_COMPUTE_KEY)
- def tearDown(self):
- self.helper.delete_instance_group_and_template()
- super().tearDown()
-
- @provide_gcp_context(GCP_COMPUTE_KEY)
- def test_run_example_dag_compute_igm(self):
- self.run_dag('example_gcp_compute_igm', CLOUD_DAG_FOLDER)
diff --git a/tests/system/providers/google/cloud/compute_igm/__init__.py b/tests/system/providers/google/cloud/compute_igm/__init__.py
@@ -0,0 +1,16 @@
+# 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.
diff --git a/...cloud/example_dags/example_compute_igm.py → .../cloud/compute_igm/example_compute_igm.py b/...cloud/example_dags/example_compute_igm.py → .../cloud/compute_igm/example_compute_igm.py
@@ -42,15 +42,17 @@
 from datetime import datetime
 
 from airflow import models
-from airflow.models.baseoperator import chain
 from airflow.providers.google.cloud.operators.compute import (
  ComputeEngineCopyInstanceTemplateOperator,
  ComputeEngineInstanceGroupUpdateManagerTemplateOperator,
 )
 
+ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID")
 GCP_PROJECT_ID = os.environ.get('GCP_PROJECT_ID', 'example-project')
 GCE_ZONE = os.environ.get('GCE_ZONE', 'europe-west1-b')
 
+DAG_ID = 'example_gcp_compute_igm'
+
 # [START howto_operator_compute_template_copy_args]
 GCE_TEMPLATE_NAME = os.environ.get('GCE_TEMPLATE_NAME', 'instance-template-test')
 GCE_NEW_TEMPLATE_NAME = os.environ.get('GCE_NEW_TEMPLATE_NAME', 'instance-template-test-new')
@@ -91,11 +93,11 @@
 
 
 with models.DAG(
- 'example_gcp_compute_igm',
+ DAG_ID,
  schedule_interval='@once', # Override to match your needs
  start_date=datetime(2021, 1, 1),
  catchup=False,
- tags=['example'],
+ tags=['example', 'igm'],
 ) as dag:
  # [START howto_operator_gce_igm_copy_template]
  gce_instance_template_copy = ComputeEngineCopyInstanceTemplateOperator(
@@ -135,9 +137,22 @@
  )
  # [END howto_operator_gce_igm_update_template_no_project_id]
 
- chain(
- gce_instance_template_copy,
- gce_instance_template_copy2,
- gce_instance_group_manager_update_template,
- gce_instance_group_manager_update_template2,
+ (
+ # TEST BODY
+ gce_instance_template_copy
+ >> gce_instance_template_copy2
+ >> gce_instance_group_manager_update_template
+ >> gce_instance_group_manager_update_template2
  )
+
+ 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()
+
+
+from tests.system.utils import get_test_run # noqa: E402
+
+# Needed to run the example DAG with pytest (see: tests/system/README.md#run_via_pytest)
+test_run = get_test_run(dag)