Fix GCS system tests (#19227)

deedmitrij · web-flow · commit eb163c81d165 · 2021-11-25T14:04:15.000+01:00
- Add '**/tmp/**' path to names of creating files
- Add separate methods and fixtures for creating\deleting required files
- Add `resource` key which disables uniform bucket level access to **GCSCreateBucketOperator**
diff --git a/airflow/providers/google/cloud/example_dags/example_gcs.py b/airflow/providers/google/cloud/example_dags/example_gcs.py
@@ -21,6 +21,7 @@
 
 import os
 from datetime import datetime
+from tempfile import gettempdir
 
 from airflow import models
 from airflow.operators.bash import BashOperator
@@ -52,10 +53,17 @@
 
 BUCKET_2 = os.environ.get("GCP_GCS_BUCKET_2", "test-gcs-example-bucket-2")
 
-PATH_TO_TRANSFORM_SCRIPT = os.environ.get('GCP_GCS_PATH_TO_TRANSFORM_SCRIPT', 'test.py')
-PATH_TO_UPLOAD_FILE = os.environ.get("GCP_GCS_PATH_TO_UPLOAD_FILE", "test-gcs-example.txt")
+temp_dir_path = gettempdir()
+PATH_TO_TRANSFORM_SCRIPT = os.environ.get(
+    "GCP_GCS_PATH_TO_TRANSFORM_SCRIPT", os.path.join(temp_dir_path, "transform_script.py")
+)
+PATH_TO_UPLOAD_FILE = os.environ.get(
+    "GCP_GCS_PATH_TO_UPLOAD_FILE", os.path.join(temp_dir_path, "test-gcs-example-upload.txt")
+)
 PATH_TO_UPLOAD_FILE_PREFIX = os.environ.get("GCP_GCS_PATH_TO_UPLOAD_FILE_PREFIX", "test-gcs-")
-PATH_TO_SAVED_FILE = os.environ.get("GCP_GCS_PATH_TO_SAVED_FILE", "test-gcs-example-download.txt")
+PATH_TO_SAVED_FILE = os.environ.get(
+    "GCP_GCS_PATH_TO_SAVED_FILE", os.path.join(temp_dir_path, "test-gcs-example-download.txt")
+)
 
 BUCKET_FILE_LOCATION = PATH_TO_UPLOAD_FILE.rpartition("/")[-1]
 
@@ -67,7 +75,16 @@
     tags=['example'],
 ) as dag:
     create_bucket1 = GCSCreateBucketOperator(
-        task_id="create_bucket1", bucket_name=BUCKET_1, project_id=PROJECT_ID
+        task_id="create_bucket1",
+        bucket_name=BUCKET_1,
+        project_id=PROJECT_ID,
+        resource={
+            "iamConfiguration": {
+                "uniformBucketLevelAccess": {
+                    "enabled": False,
+                },
+            },
+        },
     )
 
     create_bucket2 = GCSCreateBucketOperator(
@@ -179,7 +196,7 @@
     # [START howto_sensor_object_exists_task]
     gcs_object_exists = GCSObjectExistenceSensor(
         bucket=BUCKET_1,
-        object=PATH_TO_UPLOAD_FILE,
+        object=BUCKET_FILE_LOCATION,
         mode='poke',
         task_id="gcs_object_exists_task",
     )
diff --git a/tests/providers/google/cloud/operators/test_gcs_system.py b/tests/providers/google/cloud/operators/test_gcs_system.py
@@ -22,26 +22,49 @@
 from tests.test_utils.gcp_system_helpers import CLOUD_DAG_FOLDER, GoogleSystemTest, provide_gcp_context
 
 
+@pytest.fixture(scope="module")
+def helper():
+    return GcsSystemTestHelper()
+
+
+@pytest.fixture
+def file_to_upload(helper):
+    helper.create_file_to_upload()
+    yield
+    helper.remove_file_to_upload()
+
+
+@pytest.fixture
+def script_to_transform(helper):
+    helper.create_script_to_transform()
+    yield
+    helper.remove_script_to_transform()
+
+
+@pytest.fixture
+def saved_file(helper):
+    # file is created by operator inside DAG
+    yield
+    helper.remove_saved_file()
+
+
 @pytest.mark.backend("mysql", "postgres")
 @pytest.mark.credential_file(GCP_GCS_KEY)
 class GoogleCloudStorageExampleDagsTest(GoogleSystemTest):
-    helper = GcsSystemTestHelper()
-
     @provide_gcp_context(GCP_GCS_KEY)
     def setUp(self):
         super().setUp()
-        self.helper.create_test_file()
 
     @provide_gcp_context(GCP_GCS_KEY)
     def tearDown(self):
-        self.helper.remove_test_files()
-        self.helper.remove_bucket()
         super().tearDown()
 
     @provide_gcp_context(GCP_GCS_KEY)
+    @pytest.mark.usefixtures("file_to_upload", "script_to_transform", "saved_file")
     def test_run_example_dag(self):
         self.run_dag('example_gcs', CLOUD_DAG_FOLDER)
 
     @provide_gcp_context(GCP_GCS_KEY)
+    @pytest.mark.usefixtures("file_to_upload")
     def test_run_example_gcs_sensor_dag(self):
         self.run_dag('example_gcs_sensors', CLOUD_DAG_FOLDER)
diff --git a/tests/providers/google/cloud/operators/test_gcs_system_helper.py b/tests/providers/google/cloud/operators/test_gcs_system_helper.py
@@ -30,12 +30,12 @@
 
 class GcsSystemTestHelper(CommandExecutor):
     @staticmethod
-    def create_test_file():
-        # Create test file for upload
+    def create_file_to_upload():
         with open(PATH_TO_UPLOAD_FILE, "w+") as file:
             file.writelines(["This is a test file"])
 
-        # Create script for transform operator
+    @staticmethod
+    def create_script_to_transform():
         with open(PATH_TO_TRANSFORM_SCRIPT, "w+") as file:
             file.write(
                 """import sys
@@ -51,13 +51,16 @@ def create_test_file():
             )
 
     @staticmethod
-    def remove_test_files():
-        if os.path.exists(PATH_TO_UPLOAD_FILE):
-            os.remove(PATH_TO_UPLOAD_FILE)
-        if os.path.exists(PATH_TO_SAVED_FILE):
-            os.remove(PATH_TO_SAVED_FILE)
-        if os.path.exists(PATH_TO_TRANSFORM_SCRIPT):
-            os.remove(PATH_TO_TRANSFORM_SCRIPT)
+    def remove_file_to_upload():
+        os.remove(PATH_TO_UPLOAD_FILE)
+
+    @staticmethod
+    def remove_script_to_transform():
+        os.remove(PATH_TO_TRANSFORM_SCRIPT)
+
+    @staticmethod
+    def remove_saved_file():
+        os.remove(PATH_TO_SAVED_FILE)
 
     def remove_bucket(self):
         self.execute_cmd(["gsutil", "rm", "-r", f"gs://{BUCKET_1}"])
