From db38a6fc667f336d080ad48f66dc814de77c27a6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mika=C3=ABl=20Ducharme?= Date: Mon, 11 Sep 2023 05:37:44 -0400 Subject: [PATCH] GKEPodHook ignores gcp_conn_id parameter. (#34194) --- .../cloud/operators/kubernetes_engine.py | 1 + .../cloud/operators/test_kubernetes_engine.py | 37 +++++++++++++++++++ 2 files changed, 38 insertions(+) diff --git a/airflow/providers/google/cloud/operators/kubernetes_engine.py b/airflow/providers/google/cloud/operators/kubernetes_engine.py index 4f1d68c29be42..4294df224f8a7 100644 --- a/airflow/providers/google/cloud/operators/kubernetes_engine.py +++ b/airflow/providers/google/cloud/operators/kubernetes_engine.py @@ -540,6 +540,7 @@ def hook(self) -> GKEPodHook: ) hook = GKEPodHook( + gcp_conn_id=self.gcp_conn_id, cluster_url=self._cluster_url, ssl_ca_cert=self._ssl_ca_cert, ) diff --git a/tests/providers/google/cloud/operators/test_kubernetes_engine.py b/tests/providers/google/cloud/operators/test_kubernetes_engine.py index 07f0c34c7a149..328da500b68e4 100644 --- a/tests/providers/google/cloud/operators/test_kubernetes_engine.py +++ b/tests/providers/google/cloud/operators/test_kubernetes_engine.py @@ -320,6 +320,43 @@ def test_cluster_info(self, get_cluster_mock, use_internal_ip): assert cluster_url == CLUSTER_PRIVATE_URL if use_internal_ip else CLUSTER_URL assert ssl_ca_cert == SSL_CA_CERT + def test_default_gcp_conn_id(self): + gke_op = GKEStartPodOperator( + project_id=TEST_GCP_PROJECT_ID, + location=PROJECT_LOCATION, + cluster_name=CLUSTER_NAME, + task_id=PROJECT_TASK_ID, + name=TASK_NAME, + namespace=NAMESPACE, + image=IMAGE, + ) + gke_op._cluster_url = CLUSTER_URL + gke_op._ssl_ca_cert = SSL_CA_CERT + hook = gke_op.hook + + assert hook.gcp_conn_id == "google_cloud_default" + + @mock.patch( + "airflow.providers.google.common.hooks.base_google.GoogleBaseHook.get_connection", + return_value=Connection(conn_id="test_conn"), + ) + def test_gcp_conn_id(self, get_con_mock): + gke_op = GKEStartPodOperator( + project_id=TEST_GCP_PROJECT_ID, + location=PROJECT_LOCATION, + cluster_name=CLUSTER_NAME, + task_id=PROJECT_TASK_ID, + name=TASK_NAME, + namespace=NAMESPACE, + image=IMAGE, + gcp_conn_id="test_conn", + ) + gke_op._cluster_url = CLUSTER_URL + gke_op._ssl_ca_cert = SSL_CA_CERT + hook = gke_op.hook + + assert hook.gcp_conn_id == "test_conn" + @pytest.mark.parametrize( "compatible_kpo, kwargs, expected_attributes", [