specifying project id when calling wait_for_operation in delete/creat…

…e cluster (#26418) * specifying project id when calling wait_for_operation in delete/create cluster
apache · Sep 20, 2022 · e315900 · e315900
1 parent 0bca962
commit e315900
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
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):