[GitHub] [airflow] turbaszek commented on a change in pull request #7794: [AIRFLOW-7099] Improve system test for Cloud Transfer Service

[GitBox <gi...@apache.org>]

turbaszek commented on a change in pull request #7794: [AIRFLOW-7099] Improve system test for Cloud Transfer Service
URL: https://github.com/apache/airflow/pull/7794#discussion_r395997306
 
 

 ##########
 File path: tests/providers/google/cloud/operators/test_cloud_storage_transfer_service_system.py
 ##########
 @@ -17,28 +17,38 @@
 # under the License.
 import pytest
 
-from tests.providers.google.cloud.operators.test_cloud_storage_transfer_service_system_helper import (
-    GCPTransferTestHelper,
+from airflow.providers.google.cloud.example_dags.example_cloud_storage_transfer_service_gcp import (
+    GCP_PROJECT_ID, GCP_TRANSFER_FIRST_TARGET_BUCKET, GCP_TRANSFER_SECOND_TARGET_BUCKET,
 )
 from tests.providers.google.cloud.utils.gcp_authenticator import GCP_GCS_TRANSFER_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_GCS_TRANSFER_KEY)
-class GcpTransferExampleDagsSystemTest(GoogleSystemTest):
-    helper = GCPTransferTestHelper()
+@pytest.fixture
+def helper():
+    with provide_gcp_context(GCP_GCS_TRANSFER_KEY):
+        # Create buckets
+        GoogleSystemTest.create_gcs_bucket(GCP_TRANSFER_SECOND_TARGET_BUCKET, location="asia-east1")
+        GoogleSystemTest.create_gcs_bucket(GCP_TRANSFER_FIRST_TARGET_BUCKET)
+        GoogleSystemTest.upload_content_to_gcs("test_contents", GCP_TRANSFER_FIRST_TARGET_BUCKET, "test.txt")
 
-    @provide_gcp_context(GCP_GCS_TRANSFER_KEY)
-    def setUp(self):
-        super().setUp()
-        self.helper.create_gcs_buckets()
+        # Grant bucket permissions
+        project_number = GoogleSystemTest.get_project_number(GCP_PROJECT_ID)
+        account_email = f"project-{project_number}@storage-transfer-service.iam.gserviceaccount.com"
+        GoogleSystemTest.grant_bucket_access(GCP_TRANSFER_FIRST_TARGET_BUCKET, account_email)
+        GoogleSystemTest.grant_bucket_access(GCP_TRANSFER_SECOND_TARGET_BUCKET, account_email)
 
-    @provide_gcp_context(GCP_GCS_TRANSFER_KEY)
-    def tearDown(self):
-        self.helper.delete_gcs_buckets()
-        super().tearDown()
+        yield
 
+        # Remove buckets
+        GoogleSystemTest.delete_gcs_bucket(GCP_TRANSFER_SECOND_TARGET_BUCKET)
+        GoogleSystemTest.delete_gcs_bucket(GCP_TRANSFER_FIRST_TARGET_BUCKET)
 
 Review comment:
   I can, but then the example should also incorporate creation and deletion of buckets. That's, in my opinion, generates too much noise. 

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
users@infra.apache.org


With regards,
Apache Git Services