GKEPodHook ignores gcp_conn_id parameter. (#34194)

apache · Sep 11, 2023 · db38a6f · db38a6f
1 parent f70c107
commit db38a6f
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 0 deletions.
diff --git 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
@@ -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",
         [