Refactor dataproc system tests

airflow/providers/google/cloud/hooks/dataproc.py

@@ -274,7 +274,7 @@ def wait_for_operation(
  self,
  operation: Operation,
  timeout: float | None = None,
- result_retry: AsyncRetry | _MethodDefault = DEFAULT,
+ result_retry: AsyncRetry | _MethodDefault | Retry = DEFAULT,
  ) -> Any:
  """Wait for a long-lasting operation to complete."""
  try:

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

@@ -2995,7 +2995,7 @@ def __init__(
  metadata: Sequence[tuple[str, str]] = (),
  gcp_conn_id: str = "google_cloud_default",
  impersonation_chain: str | Sequence[str] | None = None,
- result_retry: AsyncRetry | _MethodDefault = DEFAULT,
+ result_retry: AsyncRetry | _MethodDefault | Retry = DEFAULT,
  asynchronous: bool = False,
  deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
  polling_interval_seconds: int = 5,

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

@@ -22,12 +22,12 @@
 import time
 from typing import TYPE_CHECKING, Sequence
 
+from google.api_core.exceptions import AlreadyExists
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.api_core.retry import Retry, exponential_sleep_generator
 from google.cloud.metastore_v1 import MetadataExport, MetadataManagementActivity
 from google.cloud.metastore_v1.types import Backup, MetadataImport, Service
 from google.cloud.metastore_v1.types.metastore import DatabaseDumpSpec, Restore
-from googleapiclient.errors import HttpError
 
 from airflow.exceptions import AirflowException
 from airflow.models import BaseOperator, BaseOperatorLink
@@ -242,9 +242,7 @@ def execute(self, context: Context) -> dict:
  )
  backup = hook.wait_for_operation(self.timeout, operation)
  self.log.info("Backup %s created successfully", self.backup_id)
- except HttpError as err:
- if err.resp.status not in (409, "409"):
- raise
+ except AlreadyExists:
  self.log.info("Backup %s already exists", self.backup_id)
  backup = hook.get_backup(
  project_id=self.project_id,
@@ -448,9 +446,7 @@ def execute(self, context: Context) -> dict:
  )
  service = hook.wait_for_operation(self.timeout, operation)
  self.log.info("Service %s created successfully", self.service_id)
- except HttpError as err:
- if err.resp.status not in (409, "409"):
- raise
+ except AlreadyExists:
  self.log.info("Instance %s already exists", self.service_id)
  service = hook.get_service(
  region=self.region,

scripts/ci/pre_commit/check_system_tests.py

@@ -35,6 +35,7 @@
 errors: list[str] = []
 
 WATCHER_APPEND_INSTRUCTION = "list(dag.tasks) >> watcher()"
+WATCHER_APPEND_INSTRUCTION_SHORT = " >> watcher()"
 
 PYTEST_FUNCTION = """
 from tests.system.utils import get_test_run # noqa: E402
@@ -52,7 +53,7 @@
 def _check_file(file: Path):
  content = file.read_text()
  if "from tests.system.utils.watcher import watcher" in content:
- index = content.find(WATCHER_APPEND_INSTRUCTION)
+ index = content.find(WATCHER_APPEND_INSTRUCTION_SHORT)
  if index == -1:
  errors.append(
  f"[red]The example {file} imports tests.system.utils.watcher "

tests/system/providers/google/cloud/dataproc/example_dataproc_batch.py

@@ -23,7 +23,7 @@
 import os
 from datetime import datetime
 
-from google.api_core.retry_async import AsyncRetry
+from google.api_core.retry import Retry
 
 from airflow.models.dag import DAG
 from airflow.providers.google.cloud.operators.dataproc import (
@@ -37,10 +37,10 @@
 from airflow.utils.trigger_rule import TriggerRule
 from tests.system.providers.google import DEFAULT_GCP_SYSTEM_TEST_PROJECT_ID
 
-ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID")
-DAG_ID = "dataproc_batch"
+ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID", "default")
 PROJECT_ID = os.environ.get("SYSTEM_TESTS_GCP_PROJECT") or DEFAULT_GCP_SYSTEM_TEST_PROJECT_ID
-REGION = "europe-west1"
+DAG_ID = "dataproc_batch"
+REGION = "europe-west3"
 
 BATCH_ID = f"batch-{ENV_ID}-{DAG_ID}".replace("_", "-")
 BATCH_ID_2 = f"batch-{ENV_ID}-{DAG_ID}-2".replace("_", "-")
@@ -77,7 +77,7 @@
  region=REGION,
  batch=BATCH_CONFIG,
  batch_id=BATCH_ID_2,
- result_retry=AsyncRetry(maximum=10.0, initial=10.0, multiplier=1.0),
+ result_retry=Retry(maximum=100.0, initial=10.0, multiplier=1.0),
  )
 
  create_batch_3 = DataprocCreateBatchOperator(
@@ -104,10 +104,6 @@
  get_batch = DataprocGetBatchOperator(
  task_id="get_batch", project_id=PROJECT_ID, region=REGION, batch_id=BATCH_ID
  )
-
- get_batch_2 = DataprocGetBatchOperator(
- task_id="get_batch_2", project_id=PROJECT_ID, region=REGION, batch_id=BATCH_ID_2
- )
  # [END how_to_cloud_dataproc_get_batch_operator]
 
  # [START how_to_cloud_dataproc_list_batches_operator]
@@ -136,6 +132,14 @@
  )
  # [END how_to_cloud_dataproc_cancel_operation_operator]
 
+ batch_cancelled_sensor = DataprocBatchSensor(
+ task_id="batch_cancelled_sensor",
+ region=REGION,
+ project_id=PROJECT_ID,
+ batch_id=BATCH_ID_4,
+ poke_interval=10,
+ )
+
  # [START how_to_cloud_dataproc_delete_batch_operator]
  delete_batch = DataprocDeleteBatchOperator(
  task_id="delete_batch", project_id=PROJECT_ID, region=REGION, batch_id=BATCH_ID
@@ -153,25 +157,30 @@
  delete_batch.trigger_rule = TriggerRule.ALL_DONE
  delete_batch_2.trigger_rule = TriggerRule.ALL_DONE
  delete_batch_3.trigger_rule = TriggerRule.ALL_DONE
- delete_batch_4.trigger_rule = TriggerRule.ALL_DONE
+ delete_batch_4.trigger_rule = TriggerRule.ALL_FAILED
 
  (
  # TEST SETUP
  [create_batch, create_batch_2, create_batch_3]
  # TEST BODY
  >> batch_async_sensor
- >> [get_batch, get_batch_2, list_batches]
+ >> get_batch
+ >> list_batches
  >> create_batch_4
  >> cancel_operation
  # TEST TEARDOWN
- >> [delete_batch, delete_batch_2, delete_batch_3, delete_batch_4]
+ >> [delete_batch, delete_batch_2, delete_batch_3]
+ >> batch_cancelled_sensor
+ >> delete_batch_4
  )
 
  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()
+
+ # Excluding sensor because we expect it to fail due to cancelled operation
+ [task for task in dag.tasks if task.task_id != "batch_cancelled_sensor"] >> watcher()
 
 
 from tests.system.utils import get_test_run # noqa: E402

tests/system/providers/google/cloud/dataproc/example_dataproc_batch_deferrable.py

@@ -37,7 +37,7 @@
 ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID")
 DAG_ID = "dataproc_batch_deferrable"
 PROJECT_ID = os.environ.get("SYSTEM_TESTS_GCP_PROJECT") or DEFAULT_GCP_SYSTEM_TEST_PROJECT_ID
-REGION = "europe-west1"
+REGION = "europe-north1"
 BATCH_ID = f"batch-{ENV_ID}-{DAG_ID}".replace("_", "-")
 BATCH_CONFIG = {
  "spark_batch": {

tests/system/providers/google/cloud/dataproc/example_dataproc_batch_persistent.py

@@ -28,17 +28,21 @@
  ClusterGenerator,
  DataprocCreateBatchOperator,
  DataprocCreateClusterOperator,
+ DataprocDeleteBatchOperator,
  DataprocDeleteClusterOperator,
 )
 from airflow.providers.google.cloud.operators.gcs import GCSCreateBucketOperator, GCSDeleteBucketOperator
 from airflow.utils.trigger_rule import TriggerRule
+from tests.system.providers.google import DEFAULT_GCP_SYSTEM_TEST_PROJECT_ID
 
-ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID")
+ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID", "default")
 DAG_ID = "dataproc_batch_ps"
-PROJECT_ID = os.environ.get("SYSTEM_TESTS_GCP_PROJECT", "default")
-BUCKET_NAME = f"bucket_{DAG_ID}_{ENV_ID}"
-REGION = "europe-west1"
-CLUSTER_NAME = f"cluster-{ENV_ID}-{DAG_ID}".replace("_", "-")
+PROJECT_ID = os.environ.get("SYSTEM_TESTS_GCP_PROJECT") or DEFAULT_GCP_SYSTEM_TEST_PROJECT_ID
+BUCKET_NAME = f"bucket_{DAG_ID}_{ENV_ID}".replace("-", "_")
+REGION = "europe-north1"
+CLUSTER_NAME_BASE = f"cluster-{DAG_ID}".replace("_", "-")
+CLUSTER_NAME_FULL = CLUSTER_NAME_BASE + f"-{ENV_ID}".replace("_", "-")
+CLUSTER_NAME = CLUSTER_NAME_BASE if len(CLUSTER_NAME_FULL) >= 33 else CLUSTER_NAME_FULL
 BATCH_ID = f"batch-{ENV_ID}-{DAG_ID}".replace("_", "-")
 
 CLUSTER_GENERATOR_CONFIG_FOR_PHS = ClusterGenerator(
@@ -98,6 +102,14 @@
  )
  # [END how_to_cloud_dataproc_create_batch_operator_with_persistent_history_server]
 
+ delete_batch = DataprocDeleteBatchOperator(
+ task_id="delete_batch",
+ project_id=PROJECT_ID,
+ region=REGION,
+ batch_id=BATCH_ID,
+ trigger_rule=TriggerRule.ALL_DONE,
+ )
+
  delete_cluster = DataprocDeleteClusterOperator(
  task_id="delete_cluster",
  project_id=PROJECT_ID,
@@ -117,6 +129,7 @@
  # TEST BODY
  >> create_batch
  # TEST TEARDOWN
+ >> delete_batch
  >> delete_cluster
  >> delete_bucket
  )

tests/system/providers/google/cloud/dataproc/example_dataproc_cluster_create_existing_stopped_cluster.py

@@ -32,14 +32,18 @@
  DataprocStopClusterOperator,
 )
 from airflow.utils.trigger_rule import TriggerRule
+from tests.system.providers.google import DEFAULT_GCP_SYSTEM_TEST_PROJECT_ID
 
-DAG_ID = "example_dataproc_cluster_create_existing_stopped_cluster"
+DAG_ID = "dataproc_create_existing_stopped_cluster"
 
-ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID")
-PROJECT_ID = os.environ.get("SYSTEMS_TESTS_GCP_PROJECTS") or ""
+ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID", "default")
+PROJECT_ID = os.environ.get("SYSTEM_TESTS_GCP_PROJECT") or DEFAULT_GCP_SYSTEM_TEST_PROJECT_ID
 
-CLUSTER_NAME = f"cluster-{ENV_ID}-{DAG_ID}".replace("_", "-")
-REGION = "europe-west1"
+
+CLUSTER_NAME_BASE = f"{DAG_ID}".replace("_", "-")
+CLUSTER_NAME_FULL = CLUSTER_NAME_BASE + f"-{ENV_ID}".replace("_", "-")
+CLUSTER_NAME = CLUSTER_NAME_BASE if len(CLUSTER_NAME_FULL) >= 33 else CLUSTER_NAME_FULL
+REGION = "europe-north1"
 
 # Cluster definition
 CLUSTER_CONFIG = {

tests/system/providers/google/cloud/dataproc/example_dataproc_cluster_deferrable.py

@@ -33,12 +33,14 @@
 from airflow.utils.trigger_rule import TriggerRule
 from tests.system.providers.google import DEFAULT_GCP_SYSTEM_TEST_PROJECT_ID
 
-ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID")
+ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID", "default")
 DAG_ID = "dataproc_cluster_def"
 PROJECT_ID = os.environ.get("SYSTEM_TESTS_GCP_PROJECT") or DEFAULT_GCP_SYSTEM_TEST_PROJECT_ID
 
-CLUSTER_NAME = f"cluster-{ENV_ID}-{DAG_ID}".replace("_", "-")
-REGION = "europe-west1"
+CLUSTER_NAME_BASE = f"cluster-{DAG_ID}".replace("_", "-")
+CLUSTER_NAME_FULL = CLUSTER_NAME_BASE + f"-{ENV_ID}".replace("_", "-")
+CLUSTER_NAME = CLUSTER_NAME_BASE if len(CLUSTER_NAME_FULL) >= 33 else CLUSTER_NAME_FULL
+REGION = "europe-north1"
 
 
 # Cluster definition

tests/system/providers/google/cloud/dataproc/example_dataproc_cluster_diagnose.py

@@ -33,12 +33,15 @@
 from airflow.utils.trigger_rule import TriggerRule
 from tests.system.providers.google import DEFAULT_GCP_SYSTEM_TEST_PROJECT_ID
 
-ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID")
+ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID", "default")
 DAG_ID = "dataproc_diagnose_cluster"
 PROJECT_ID = os.environ.get("SYSTEM_TESTS_GCP_PROJECT") or DEFAULT_GCP_SYSTEM_TEST_PROJECT_ID
 
-CLUSTER_NAME = f"cluster-{ENV_ID}-{DAG_ID}".replace("_", "-")
-REGION = "europe-west1"
+CLUSTER_NAME_BASE = f"cluster-{DAG_ID}".replace("_", "-")
+CLUSTER_NAME_FULL = CLUSTER_NAME_BASE + f"-{ENV_ID}".replace("_", "-")
+CLUSTER_NAME = CLUSTER_NAME_BASE if len(CLUSTER_NAME_FULL) >= 33 else CLUSTER_NAME_FULL
+
+REGION = "europe-north1"
 
 
 # Cluster definition
@@ -102,7 +105,7 @@
  # TEST SETUP
  create_cluster
  # TEST BODY
- >> diagnose_cluster
+ >> [diagnose_cluster, diagnose_cluster_deferrable]
  # TEST TEARDOWN
  >> delete_cluster
  )

tests/system/providers/google/cloud/dataproc/example_dataproc_cluster_generator.py

@@ -24,33 +24,39 @@
 
 import os
 from datetime import datetime
-from pathlib import Path
 
 from airflow.models.dag import DAG
 from airflow.providers.google.cloud.operators.dataproc import (
  ClusterGenerator,
  DataprocCreateClusterOperator,
  DataprocDeleteClusterOperator,
 )
-from airflow.providers.google.cloud.operators.gcs import GCSCreateBucketOperator, GCSDeleteBucketOperator
-from airflow.providers.google.cloud.transfers.local_to_gcs import LocalFilesystemToGCSOperator
+from airflow.providers.google.cloud.operators.gcs import (
+ GCSCreateBucketOperator,
+ GCSDeleteBucketOperator,
+ GCSSynchronizeBucketsOperator,
+)
 from airflow.utils.trigger_rule import TriggerRule
+from tests.system.providers.google import DEFAULT_GCP_SYSTEM_TEST_PROJECT_ID
 
-ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID")
+ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID", "default")
 DAG_ID = "dataproc_cluster_generation"
-PROJECT_ID = os.environ.get("SYSTEM_TESTS_GCP_PROJECT", "default")
+PROJECT_ID = os.environ.get("SYSTEM_TESTS_GCP_PROJECT") or DEFAULT_GCP_SYSTEM_TEST_PROJECT_ID
 
 BUCKET_NAME = f"bucket_{DAG_ID}_{ENV_ID}"
-CLUSTER_NAME = f"cluster-{ENV_ID}-{DAG_ID}".replace("_", "-")
-REGION = "europe-west1"
-ZONE = "europe-west1-b"
-INIT_FILE_SRC = str(Path(__file__).parent / "resources" / "pip-install.sh")
+RESOURCE_DATA_BUCKET = "airflow-system-tests-resources"
+INIT_FILE = "pip-install.sh"
+GCS_INIT_FILE = f"gs://{RESOURCE_DATA_BUCKET}/dataproc/{INIT_FILE}"
 
-# Cluster definition: Generating Cluster Config for DataprocCreateClusterOperator
-# [START how_to_cloud_dataproc_create_cluster_generate_cluster_config]
+CLUSTER_NAME_BASE = f"cluster-{DAG_ID}".replace("_", "-")
+CLUSTER_NAME_FULL = CLUSTER_NAME_BASE + f"-{ENV_ID}".replace("_", "-")
+CLUSTER_NAME = CLUSTER_NAME_BASE if len(CLUSTER_NAME_FULL) >= 33 else CLUSTER_NAME_FULL
 
-INIT_FILE = "pip-install.sh"
+REGION = "us-east4"
+ZONE = "us-east4-a"
 
+# Cluster definition: Generating Cluster Config for DataprocCreateClusterOperator
+# [START how_to_cloud_dataproc_create_cluster_generate_cluster_config]
 CLUSTER_GENERATOR_CONFIG = ClusterGenerator(
  project_id=PROJECT_ID,
  zone=ZONE,
@@ -60,7 +66,7 @@
  worker_disk_size=32,
  num_workers=2,
  storage_bucket=BUCKET_NAME,
- init_actions_uris=[f"gs://{BUCKET_NAME}/{INIT_FILE}"],
+ init_actions_uris=[GCS_INIT_FILE],
  metadata={"PIP_PACKAGES": "pyyaml requests pandas openpyxl"},
  num_preemptible_workers=1,
  preemptibility="PREEMPTIBLE",
@@ -80,11 +86,13 @@
  task_id="create_bucket", bucket_name=BUCKET_NAME, project_id=PROJECT_ID
  )
 
- upload_file = LocalFilesystemToGCSOperator(
- task_id="upload_file",
- src=INIT_FILE_SRC,
- dst=INIT_FILE,
- bucket=BUCKET_NAME,
+ move_init_file = GCSSynchronizeBucketsOperator(
+ task_id="move_init_file",
+ source_bucket=RESOURCE_DATA_BUCKET,
+ source_object="dataproc",
+ destination_bucket=BUCKET_NAME,
+ destination_object="dataproc",
+ recursive=True,
  )
 
  # [START how_to_cloud_dataproc_create_cluster_generate_cluster_config_operator]
@@ -114,7 +122,7 @@
  (
  # TEST SETUP
  create_bucket
- >> upload_file
+ >> move_init_file
  # TEST BODY
  >> create_dataproc_cluster
  # TEST TEARDOWN