Translate system tests migration (AIP-47) (#25340)

airflow/providers/google/cloud/operators/translate_speech.py

@@ -25,6 +25,7 @@
 from airflow.models import BaseOperator
 from airflow.providers.google.cloud.hooks.speech_to_text import CloudSpeechToTextHook
 from airflow.providers.google.cloud.hooks.translate import CloudTranslateHook
+from airflow.providers.google.common.links.storage import FileDetailsLink
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
@@ -109,6 +110,7 @@ class CloudTranslateSpeechOperator(BaseOperator):
         'gcp_conn_id',
         'impersonation_chain',
     )
+    operator_extra_links = (FileDetailsLink(),)
     # [END translate_speech_template_fields]
 
     def __init__(
@@ -172,6 +174,12 @@ def execute(self, context: 'Context') -> dict:
                 model=self.model,
             )
             self.log.info('Translated output: %s', translation)
+            FileDetailsLink.persist(
+                context=context,
+                task_instance=self,
+                uri=self.audio["uri"][5:],
+                project_id=self.project_id or translate_hook.project_id,
+            )
             return translation
         except ValueError as e:
             self.log.error('An error has been thrown from translate speech method:')

docs/apache-airflow-providers-google/operators/cloud/translate.rst

@@ -40,7 +40,7 @@ Using the operator
 
 Basic usage of the operator:
 
-.. exampleinclude:: /../../airflow/providers/google/cloud/example_dags/example_translate.py
+.. exampleinclude:: /../../tests/system/providers/google/cloud/translate/example_translate.py
       :language: python
       :dedent: 4
       :start-after: [START howto_operator_translate_text]
@@ -49,7 +49,7 @@ Basic usage of the operator:
 The result of translation is available as dictionary or array of dictionaries accessible via the usual
 XCom mechanisms of Airflow:
 
-.. exampleinclude:: /../../airflow/providers/google/cloud/example_dags/example_translate.py
+.. exampleinclude:: /../../tests/system/providers/google/cloud/translate/example_translate.py
       :language: python
       :dedent: 4
       :start-after: [START howto_operator_translate_access]

docs/apache-airflow-providers-google/operators/cloud/translate_speech.rst

@@ -43,7 +43,7 @@ for more information, see: https://googleapis.github.io/google-cloud-python/late
 
 Arguments for translation need to be specified.
 
-.. exampleinclude:: /../../airflow/providers/google/cloud/example_dags/example_translate_speech.py
+.. exampleinclude:: /../../tests/system/providers/google/cloud/translate_speech/example_translate_speech.py
       :language: python
       :start-after: [START howto_operator_translate_speech_arguments]
       :end-before: [END howto_operator_translate_speech_arguments]
@@ -52,7 +52,7 @@ Arguments for translation need to be specified.
 Using the operator
 """"""""""""""""""
 
-.. exampleinclude:: /../../airflow/providers/google/cloud/example_dags/example_translate_speech.py
+.. exampleinclude:: /../../tests/system/providers/google/cloud/translate_speech/example_translate_speech.py
       :language: python
       :dedent: 4
       :start-after: [START howto_operator_translate_speech]

tests/providers/google/cloud/operators/test_translate.py

@@ -46,7 +46,8 @@ def test_minimal_green_path(self, mock_hook):
             task_id='id',
             impersonation_chain=IMPERSONATION_CHAIN,
         )
-        return_value = op.execute(context=None)
+        context = mock.MagicMock()
+        return_value = op.execute(context=context)
         mock_hook.assert_called_once_with(
             gcp_conn_id=GCP_CONN_ID,
             impersonation_chain=IMPERSONATION_CHAIN,

tests/providers/google/cloud/operators/test_translate_speech.py

@@ -64,7 +64,8 @@ def test_minimal_green_path(self, mock_translate_hook, mock_speech_hook):
             task_id='id',
             impersonation_chain=IMPERSONATION_CHAIN,
         )
-        return_value = op.execute(context=None)
+        context = mock.MagicMock()
+        return_value = op.execute(context=context)
 
         mock_speech_hook.assert_called_once_with(
             gcp_conn_id=GCP_CONN_ID,

tests/providers/google/cloud/operators/test_translate_system.py → tests/system/providers/google/cloud/translate/__init__.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
@@ -15,21 +14,3 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-import pytest
-
-from tests.providers.google.cloud.utils.gcp_authenticator import GCP_AI_KEY
-from tests.test_utils.gcp_system_helpers import CLOUD_DAG_FOLDER, GoogleSystemTest, provide_gcp_context
-
-
-@pytest.mark.backend("mysql", "postgres")
-@pytest.mark.credential_file(GCP_AI_KEY)
-class CloudTranslateExampleDagsSystemTest(GoogleSystemTest):
-    def setUp(self):
-        super().setUp()
-
-    @provide_gcp_context(GCP_AI_KEY)
-    def test_run_example_dag_function(self):
-        self.run_dag('example_gcp_translate', CLOUD_DAG_FOLDER)
-
-    def tearDown(self):
-        super().tearDown()

airflow/providers/google/cloud/example_dags/example_translate.py → tests/system/providers/google/cloud/translate/example_translate.py

@@ -27,26 +27,42 @@
 from airflow.operators.bash import BashOperator
 from airflow.providers.google.cloud.operators.translate import CloudTranslateTextOperator
 
+DAG_ID = "example_gcp_translate"
+
 with models.DAG(
-    'example_gcp_translate',
-    schedule_interval='@once',  # Override to match your needs
+    DAG_ID,
+    schedule_interval="@once",  # Override to match your needs
     start_date=datetime(2021, 1, 1),
     catchup=False,
-    tags=['example'],
+    tags=["example"],
 ) as dag:
     # [START howto_operator_translate_text]
     product_set_create = CloudTranslateTextOperator(
-        task_id='translate',
-        values=['zażółć gęślą jaźń'],
-        target_language='en',
-        format_='text',
+        task_id="translate",
+        values=["zażółć gęślą jaźń"],
+        target_language="en",
+        format_="text",
         source_language=None,
-        model='base',
+        model="base",
     )
     # [END howto_operator_translate_text]
     # [START howto_operator_translate_access]
     translation_access = BashOperator(
-        task_id='access', bash_command="echo '{{ task_instance.xcom_pull(\"translate\")[0] }}'"
+        task_id="access", bash_command="echo '{{ task_instance.xcom_pull(\"translate\")[0] }}'"
     )
-    product_set_create >> translation_access
     # [END howto_operator_translate_access]
+    product_set_create >> translation_access
+
+    # ### Everything below this line is not part of example ###
+    # ### Just for system tests purpose ###
+    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)

tests/system/providers/google/cloud/translate_speech/__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.

airflow/providers/google/cloud/example_dags/example_translate_speech.py → tests/system/providers/google/cloud/translate_speech/example_translate_speech.py

@@ -20,14 +20,20 @@
 from datetime import datetime
 
 from airflow import models
+from airflow.providers.google.cloud.operators.gcs import GCSCreateBucketOperator, GCSDeleteBucketOperator
 from airflow.providers.google.cloud.operators.text_to_speech import CloudTextToSpeechSynthesizeOperator
 from airflow.providers.google.cloud.operators.translate_speech import CloudTranslateSpeechOperator
+from airflow.utils.trigger_rule import TriggerRule
 
-GCP_PROJECT_ID = os.environ.get("GCP_PROJECT_ID", "example-project")
-BUCKET_NAME = os.environ.get("GCP_TRANSLATE_SPEECH_TEST_BUCKET", "INVALID BUCKET NAME")
+ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID")
+PROJECT_ID = os.environ.get("SYSTEM_TESTS_GCP_PROJECT")
+
+DAG_ID = "example_gcp_translate_speech"
+
+BUCKET_NAME = f"bucket_{DAG_ID}_{ENV_ID}"
 
 # [START howto_operator_translate_speech_gcp_filename]
-FILENAME = "gcp-speech-test-file"
+FILE_NAME = f"test-translate-speech-file-{DAG_ID}-{ENV_ID}"
 # [END howto_operator_translate_speech_gcp_filename]
 
 # [START howto_operator_text_to_speech_api_arguments]
@@ -38,7 +44,7 @@
 
 # [START howto_operator_translate_speech_arguments]
 CONFIG = {"encoding": "LINEAR16", "language_code": "en_US"}
-AUDIO = {"uri": f"gs://{BUCKET_NAME}/{FILENAME}"}
+AUDIO = {"uri": f"gs://{BUCKET_NAME}/{FILE_NAME}"}
 TARGET_LANGUAGE = 'pl'
 FORMAT = 'text'
 MODEL = 'base'
@@ -47,24 +53,26 @@
 
 
 with models.DAG(
-    "example_gcp_translate_speech",
+    DAG_ID,
     schedule_interval='@once',  # Override to match your needs
     start_date=datetime(2021, 1, 1),
     catchup=False,
-    tags=['example'],
+    tags=["example"],
 ) as dag:
+    create_bucket = GCSCreateBucketOperator(task_id="create_bucket", bucket_name=BUCKET_NAME)
+
     text_to_speech_synthesize_task = CloudTextToSpeechSynthesizeOperator(
-        project_id=GCP_PROJECT_ID,
+        project_id=PROJECT_ID,
         input_data=INPUT,
         voice=VOICE,
         audio_config=AUDIO_CONFIG,
         target_bucket_name=BUCKET_NAME,
-        target_filename=FILENAME,
+        target_filename=FILE_NAME,
         task_id="text_to_speech_synthesize_task",
     )
     # [START howto_operator_translate_speech]
     translate_speech_task = CloudTranslateSpeechOperator(
-        project_id=GCP_PROJECT_ID,
+        project_id=PROJECT_ID,
         audio=AUDIO,
         config=CONFIG,
         target_language=TARGET_LANGUAGE,
@@ -83,4 +91,30 @@
         task_id='translate_speech_task2',
     )
     # [END howto_operator_translate_speech]
-    text_to_speech_synthesize_task >> translate_speech_task >> translate_speech_task2
+
+    delete_bucket = GCSDeleteBucketOperator(
+        task_id="delete_bucket", bucket_name=BUCKET_NAME, trigger_rule=TriggerRule.ALL_DONE
+    )
+
+    (
+        # TEST SETUP
+        create_bucket
+        # TEST BODY
+        >> text_to_speech_synthesize_task
+        >> translate_speech_task
+        >> translate_speech_task2
+        # TEST TEARDOWN
+        >> delete_bucket
+    )
+
+    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)