Fix GCS system tests (#19227)

- 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**
apache · Nov 25, 2021 · eb163c8 · eb163c8
1 parent 744d11b
commit eb163c8
Show file tree

Hide file tree

Showing 3 changed files with 63 additions and 20 deletions.
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}"])