specifying project id when calling wait_for_operation in delete/create cluster (#26418)

SPTKL · web-flow · commit e31590039634 · 2022-09-20T09:46:16.000+02:00
* specifying project id when calling wait_for_operation in delete/create cluster
diff --git a/airflow/providers/google/cloud/hooks/kubernetes_engine.py b/airflow/providers/google/cloud/hooks/kubernetes_engine.py
@@ -180,7 +180,7 @@ def delete_cluster(
                 retry=retry,
                 timeout=timeout,
             )
-            resource = self.wait_for_operation(resource)
+            resource = self.wait_for_operation(resource, project_id)
             # Returns server-defined url for the resource
             return resource.self_link
         except NotFound as error:
@@ -234,7 +234,7 @@ def create_cluster(
                 retry=retry,
                 timeout=timeout,
             )
-            resource = self.wait_for_operation(resource)
+            resource = self.wait_for_operation(resource, project_id)
 
             return resource.target_link
         except AlreadyExists as error:
diff --git a/tests/providers/google/cloud/hooks/test_kubernetes_engine.py b/tests/providers/google/cloud/hooks/test_kubernetes_engine.py
@@ -73,7 +73,7 @@ def test_delete_cluster(self, wait_mock):
             retry=retry_mock,
             timeout=timeout_mock,
         )
-        wait_mock.assert_called_once_with(client_delete.return_value)
+        wait_mock.assert_called_once_with(client_delete.return_value, TEST_GCP_PROJECT_ID)
 
     @mock.patch(GKE_STRING.format("GKEHook.log"))
     @mock.patch(GKE_STRING.format("GKEHook.wait_for_operation"))
@@ -130,7 +130,7 @@ def test_create_cluster_proto(self, wait_mock):
             retry=retry_mock,
             timeout=timeout_mock,
         )
-        wait_mock.assert_called_once_with(client_create.return_value)
+        wait_mock.assert_called_once_with(client_create.return_value, TEST_GCP_PROJECT_ID)
 
     @mock.patch(GKE_STRING.format("Cluster.from_json"))
     @mock.patch(GKE_STRING.format("GKEHook.wait_for_operation"))
@@ -151,7 +151,7 @@ def test_create_cluster_dict(self, wait_mock, convert_mock):
             retry=retry_mock,
             timeout=timeout_mock,
         )
-        wait_mock.assert_called_once_with(client_create.return_value)
+        wait_mock.assert_called_once_with(client_create.return_value, TEST_GCP_PROJECT_ID)
 
     @mock.patch(GKE_STRING.format("GKEHook.wait_for_operation"))
     def test_create_cluster_error(self, wait_mock):
