Create system test for K8s and dataproc operators

Author: MaksYermak

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,

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. "

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,

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]

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,

tests/providers/google/cloud/operators/test_kubernetes_engine_system.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)

tests/system/providers/google/dataproc/example_dataproc_gke.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)