Create system test for K8s and dataproc operators

apache · Apr 25, 2022 · c46d04e · c46d04e
1 parent 35cbc89
commit c46d04e
Show file tree

Hide file tree

Showing 8 changed files with 239 additions and 47 deletions.
diff --git a/airflow/providers/google/cloud/hooks/dataproc.py b/airflow/providers/google/cloud/hooks/dataproc.py
@@ -294,7 +294,7 @@ def create_cluster(
  region: str,
  project_id: str,
  cluster_name: str,
- cluster_config: Union[Dict, Cluster],
+ cluster_config: Union[Dict, Cluster, None],
  virtual_cluster_config: Optional[Dict] = None,
  run_in_gke_cluster: Optional[bool] = False,
  labels: Optional[Dict[str, str]] = None,

diff --git a/airflow/providers/google/cloud/operators/dataproc.py b/airflow/providers/google/cloud/operators/dataproc.py
@@ -459,7 +459,7 @@ def __init__(
  cluster_name: str,
  region: Optional[str] = None,
  project_id: Optional[str] = None,
- cluster_config: Optional[Dict] = None,
+ cluster_config: Optional[Union[Dict, Cluster]] = None,
  virtual_cluster_config: Optional[Dict] = None,
  run_in_gke_cluster: bool = False,
  labels: Optional[Dict] = None,
@@ -482,7 +482,7 @@ def __init__(
  region = 'global'
 
  # TODO: remove one day
- if cluster_config is None:
+ if cluster_config is None and not run_in_gke_cluster:
  warnings.warn(
  f"Passing cluster parameters by keywords to `{type(self).__name__}` will be deprecated. "
  "Please provide cluster_config object using `cluster_config` parameter. "

diff --git a/docs/apache-airflow-providers-google/operators/cloud/dataproc.rst b/docs/apache-airflow-providers-google/operators/cloud/dataproc.rst
@@ -58,6 +58,24 @@ With this configuration we can create the cluster:
  :start-after: [START how_to_cloud_dataproc_create_cluster_operator]
  :end-before: [END how_to_cloud_dataproc_create_cluster_operator]
 
+For create Dataproc cluster in Google Kubernetes Engine you should enable run_in_gke_cluster flag
+and use this cluster configuration:
+
+.. exampleinclude:: /../../tests/system/providers/google/dataproc/example_dataproc_gke.py
+ :language: python
+ :dedent: 0
+ :start-after: [START how_to_cloud_dataproc_create_cluster_in_gke_config]
+ :end-before: [END how_to_cloud_dataproc_create_cluster_in_gke_config]
+
+With this configuration we can create the cluster:
+:class:`~airflow.providers.google.cloud.operators.dataproc.DataprocCreateClusterOperator`
+
+.. exampleinclude:: /../../tests/system/providers/google/dataproc/example_dataproc_gke.py
+ :language: python
+ :dedent: 4
+ :start-after: [START how_to_cloud_dataproc_create_cluster_operator_in_gke]
+ :end-before: [END how_to_cloud_dataproc_create_cluster_operator_in_gke]
+
 Generating Cluster Config
 ^^^^^^^^^^^^^^^^^^^^^^^^^
 You can also generate **CLUSTER_CONFIG** using functional API,

diff --git a/docs/apache-airflow-providers-google/operators/cloud/kubernetes_engine.rst b/docs/apache-airflow-providers-google/operators/cloud/kubernetes_engine.rst
@@ -43,7 +43,7 @@ Create GKE cluster
 
 Here is an example of a cluster definition:
 
-.. exampleinclude:: /../../airflow/providers/google/cloud/example_dags/example_kubernetes_engine.py
+.. exampleinclude:: /../../tests/system/providers/google/kubernetes_engine/example_kubernetes_engine.py
  :language: python
  :start-after: [START howto_operator_gcp_gke_create_cluster_definition]
  :end-before: [END howto_operator_gcp_gke_create_cluster_definition]
@@ -53,7 +53,7 @@ A dict object like this, or a
 definition, is required when creating a cluster with
 :class:`~airflow.providers.google.cloud.operators.kubernetes_engine.GKECreateClusterOperator`.
 
-.. exampleinclude:: /../../airflow/providers/google/cloud/example_dags/example_kubernetes_engine.py
+.. exampleinclude:: /../../tests/system/providers/google/kubernetes_engine/example_kubernetes_engine.py
  :language: python
  :dedent: 4
  :start-after: [START howto_operator_gke_create_cluster]
@@ -68,7 +68,7 @@ To delete a cluster, use
 :class:`~airflow.providers.google.cloud.operators.kubernetes_engine.GKEDeleteClusterOperator`.
 This would also delete all the nodes allocated to the cluster.
 
-.. exampleinclude:: /../../airflow/providers/google/cloud/example_dags/example_kubernetes_engine.py
+.. exampleinclude:: /../../tests/system/providers/google/kubernetes_engine/example_kubernetes_engine.py
  :language: python
  :dedent: 4
  :start-after: [START howto_operator_gke_delete_cluster]
@@ -117,15 +117,15 @@ is the path ``/airflow/xcom``. To provide values to the XCom, ensure your Pod wr
 ``return.json`` in the sidecar. The contents of this can then be used downstream in your DAG.
 Here is an example of it being used:
 
-.. exampleinclude:: /../../airflow/providers/google/cloud/example_dags/example_kubernetes_engine.py
+.. exampleinclude:: /../../tests/system/providers/google/kubernetes_engine/example_kubernetes_engine.py
  :language: python
  :dedent: 4
  :start-after: [START howto_operator_gke_start_pod_xcom]
  :end-before: [END howto_operator_gke_start_pod_xcom]
 
 And then use it in other operators:
 
-.. exampleinclude:: /../../airflow/providers/google/cloud/example_dags/example_kubernetes_engine.py
+.. exampleinclude:: /../../tests/system/providers/google/kubernetes_engine/example_kubernetes_engine.py
  :language: python
  :dedent: 4
  :start-after: [START howto_operator_gke_xcom_result]

diff --git a/tests/providers/google/cloud/operators/test_dataproc.py b/tests/providers/google/cloud/operators/test_dataproc.py
@@ -113,6 +113,21 @@
  ],
  "endpoint_config": {},
 }
+VIRTUAL_CLUSTER_CONFIG = {
+ "kubernetes_cluster_config": {
+ "gke_cluster_config": {
+ "gke_cluster_target": "projects/project_id/locations/region/clusters/gke_cluster_name",
+ "node_pool_target": [
+ {
+ "node_pool": "projects/project_id/locations/region/clusters/gke_cluster_name/nodePools/dp", # noqa
+ "roles": ["DEFAULT"],
+ }
+ ],
+ },
+ "kubernetes_software_config": {"component_version": {"SPARK": b'3'}},
+ },
+ "staging_bucket": "test-staging-bucket",
+}
 
 CONFIG_WITH_CUSTOM_IMAGE_FAMILY = {
  "gce_cluster_config": {
@@ -424,6 +439,8 @@ def test_execute(self, mock_hook, to_dict_mock):
  'metadata': METADATA,
  'cluster_config': CONFIG,
  'labels': LABELS,
+ 'run_in_gke_cluster': False,
+ 'virtual_cluster_config': None,
  }
  expected_calls = self.extra_links_expected_calls_base + [
  call.hook().create_cluster(**create_cluster_args),
@@ -457,6 +474,57 @@ def test_execute(self, mock_hook, to_dict_mock):
  execution_date=None,
  )
 
+ @mock.patch(DATAPROC_PATH.format("Cluster.to_dict"))
+ @mock.patch(DATAPROC_PATH.format("DataprocHook"))
+ def test_execute_in_gke(self, mock_hook, to_dict_mock):
+ self.extra_links_manager_mock.attach_mock(mock_hook, 'hook')
+ mock_hook.return_value.create_cluster.result.return_value = None
+ create_cluster_args = {
+ 'region': GCP_LOCATION,
+ 'project_id': GCP_PROJECT,
+ 'cluster_name': CLUSTER_NAME,
+ 'request_id': REQUEST_ID,
+ 'retry': RETRY,
+ 'timeout': TIMEOUT,
+ 'metadata': METADATA,
+ 'cluster_config': None,
+ 'labels': LABELS,
+ 'run_in_gke_cluster': True,
+ 'virtual_cluster_config': VIRTUAL_CLUSTER_CONFIG,
+ }
+ expected_calls = self.extra_links_expected_calls_base + [
+ call.hook().create_cluster(**create_cluster_args),
+ ]
+
+ op = DataprocCreateClusterOperator(
+ task_id=TASK_ID,
+ region=GCP_LOCATION,
+ labels=LABELS,
+ cluster_name=CLUSTER_NAME,
+ project_id=GCP_PROJECT,
+ run_in_gke_cluster=True,
+ virtual_cluster_config=VIRTUAL_CLUSTER_CONFIG,
+ request_id=REQUEST_ID,
+ gcp_conn_id=GCP_CONN_ID,
+ retry=RETRY,
+ timeout=TIMEOUT,
+ metadata=METADATA,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ )
+ op.execute(context=self.mock_context)
+ mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN)
+ mock_hook.return_value.create_cluster.assert_called_once_with(**create_cluster_args)
+
+ # Test whether xcom push occurs before create cluster is called
+ self.extra_links_manager_mock.assert_has_calls(expected_calls, any_order=False)
+
+ to_dict_mock.assert_called_once_with(mock_hook().create_cluster().result())
+ self.mock_ti.xcom_push.assert_called_once_with(
+ key="conf",
+ value=DATAPROC_CLUSTER_CONF_EXPECTED,
+ execution_date=None,
+ )
+
  @mock.patch(DATAPROC_PATH.format("Cluster.to_dict"))
  @mock.patch(DATAPROC_PATH.format("DataprocHook"))
  def test_execute_if_cluster_exists(self, mock_hook, to_dict_mock):
@@ -488,6 +556,8 @@ def test_execute_if_cluster_exists(self, mock_hook, to_dict_mock):
  retry=RETRY,
  timeout=TIMEOUT,
  metadata=METADATA,
+ run_in_gke_cluster=False,
+ virtual_cluster_config=None,
  )
  mock_hook.return_value.get_cluster.assert_called_once_with(
  region=GCP_LOCATION,

diff --git a/tests/providers/google/cloud/operators/test_kubernetes_engine_system.py b/tests/providers/google/cloud/operators/test_kubernetes_engine_system.py
diff --git a/tests/system/providers/google/dataproc/example_dataproc_gke.py b/tests/system/providers/google/dataproc/example_dataproc_gke.py
@@ -0,0 +1,124 @@
+#
+# 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.
+"""
+Example Airflow DAG that show how to create a Dataproc cluster in Google Kubernetes Engine.
+"""
+
+import os
+from datetime import datetime
+
+from airflow import models
+from airflow.providers.google.cloud.operators.dataproc import (
+ DataprocCreateClusterOperator,
+ DataprocDeleteClusterOperator,
+)
+from airflow.providers.google.cloud.operators.kubernetes_engine import (
+ GKECreateClusterOperator,
+ GKEDeleteClusterOperator,
+)
+from airflow.utils.trigger_rule import TriggerRule
+
+ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID")
+DAG_ID = "dataproc-gke"
+PROJECT_ID = os.environ.get("SYSTEM_TESTS_GCP_PROJECT")
+
+REGION = "us-central1"
+CLUSTER_NAME = f"cluster-test-build-in-gke{ENV_ID}"
+GKE_CLUSTER_NAME = f"test-dataproc-gke-cluster-{ENV_ID}"
+GKE_CLUSTER_CONFIG = {
+ "name": GKE_CLUSTER_NAME,
+ "workload_identity_config": {
+ "workload_pool": f"{PROJECT_ID}.svc.id.goog",
+ },
+ "initial_node_count": 1,
+}
+
+# [START how_to_cloud_dataproc_create_cluster_in_gke_config]
+
+VIRTUAL_CLUSTER_CONFIG = {
+ "kubernetes_cluster_config": {
+ "gke_cluster_config": {
+ "gke_cluster_target": f"projects/{PROJECT_ID}/locations/{REGION}/clusters/{GKE_CLUSTER_NAME}",
+ "node_pool_target": [
+ {
+ "node_pool": f"projects/{PROJECT_ID}/locations/{REGION}/clusters/{GKE_CLUSTER_NAME}/nodePools/dp", # noqa
+ "roles": ["DEFAULT"],
+ }
+ ],
+ },
+ "kubernetes_software_config": {"component_version": {"SPARK": b'3'}},
+ },
+ "staging_bucket": "test-staging-bucket",
+}
+
+# [END how_to_cloud_dataproc_create_cluster_in_gke_config]
+
+
+with models.DAG(
+ DAG_ID,
+ schedule_interval='@once',
+ start_date=datetime(2021, 1, 1),
+ catchup=False,
+ tags=['example'],
+) as dag:
+ create_gke_cluster = GKECreateClusterOperator(
+ task_id="create_gke_cluster",
+ project_id=PROJECT_ID,
+ location=REGION,
+ body=GKE_CLUSTER_CONFIG,
+ )
+
+ # [START how_to_cloud_dataproc_create_cluster_operator_in_gke]
+ create_cluster_in_gke = DataprocCreateClusterOperator(
+ task_id="create_cluster_in_gke",
+ project_id=PROJECT_ID,
+ region=REGION,
+ cluster_name=CLUSTER_NAME,
+ run_in_gke_cluster=True,
+ virtual_cluster_config=VIRTUAL_CLUSTER_CONFIG,
+ )
+ # [END how_to_cloud_dataproc_create_cluster_operator_in_gke]
+
+ delete_dataproc_cluster = DataprocDeleteClusterOperator(
+ task_id="delete_dataproc_cluster",
+ project_id=PROJECT_ID,
+ cluster_name=CLUSTER_NAME,
+ region=REGION,
+ trigger_rule=TriggerRule.ALL_DONE,
+ )
+
+ delete_gke_cluster = GKEDeleteClusterOperator(
+ task_id="delete_gke_cluster",
+ name=GKE_CLUSTER_NAME,
+ project_id=PROJECT_ID,
+ location=REGION,
+ trigger_rule=TriggerRule.ALL_DONE,
+ )
+
+ create_gke_cluster >> create_cluster_in_gke >> [delete_dataproc_cluster, delete_gke_cluster]
+
+ 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)
diff --git a/...example_dags/example_kubernetes_engine.py → ...netes_engine/example_kubernetes_engine.py b/...example_dags/example_kubernetes_engine.py → ...netes_engine/example_kubernetes_engine.py
@@ -30,16 +30,19 @@
  GKEStartPodOperator,
 )
 
-GCP_PROJECT_ID = os.environ.get("GCP_PROJECT_ID", "example-project")
-GCP_LOCATION = os.environ.get("GCP_GKE_LOCATION", "europe-north1-a")
-CLUSTER_NAME = os.environ.get("GCP_GKE_CLUSTER_NAME", "cluster-name")
+ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID")
+DAG_ID = "kubernetes_engine"
+GCP_PROJECT_ID = os.environ.get("SYSTEM_TESTS_GCP_PROJECT")
+
+GCP_LOCATION = "europe-north1-a"
+CLUSTER_NAME = f"cluster-name-test-build-{ENV_ID}"
 
 # [START howto_operator_gcp_gke_create_cluster_definition]
 CLUSTER = {"name": CLUSTER_NAME, "initial_node_count": 1}
 # [END howto_operator_gcp_gke_create_cluster_definition]
 
 with models.DAG(
- "example_gcp_gke",
+ DAG_ID,
  schedule_interval='@once', # Override to match your needs
  start_date=datetime(2021, 1, 1),
  catchup=False,
@@ -101,3 +104,15 @@
  create_cluster >> pod_task >> delete_cluster
  create_cluster >> pod_task_xcom >> delete_cluster
  pod_task_xcom >> pod_task_xcom_result
+
+ 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)