[GitHub] feng-tao closed pull request #3711: [AIRFLOW-2863] Fix GKEClusterHook catching wrong exception

[GitBox <gi...@apache.org>]

feng-tao closed pull request #3711: [AIRFLOW-2863] Fix GKEClusterHook catching wrong exception
URL: https://github.com/apache/incubator-airflow/pull/3711
 
 
   

This is a PR merged from a forked repository.
As GitHub hides the original diff on merge, it is displayed below for
the sake of provenance:

As this is a foreign pull request (from a fork), the diff is supplied
below (as it won't show otherwise due to GitHub magic):

diff --git a/airflow/contrib/hooks/gcp_container_hook.py b/airflow/contrib/hooks/gcp_container_hook.py
index d36d796d76..227cd3d124 100644
--- a/airflow/contrib/hooks/gcp_container_hook.py
+++ b/airflow/contrib/hooks/gcp_container_hook.py
@@ -23,7 +23,7 @@
 from airflow import AirflowException, version
 from airflow.hooks.base_hook import BaseHook
 
-from google.api_core.exceptions import AlreadyExists
+from google.api_core.exceptions import AlreadyExists, NotFound
 from google.api_core.gapic_v1.method import DEFAULT
 from google.cloud import container_v1, exceptions
 from google.cloud.container_v1.gapic.enums import Operation
@@ -141,7 +141,7 @@ def delete_cluster(self, name, retry=DEFAULT, timeout=DEFAULT):
             op = self.wait_for_operation(op)
             # Returns server-defined url for the resource
             return op.self_link
-        except exceptions.NotFound as error:
+        except NotFound as error:
             self.log.info('Assuming Success: ' + error.message)
 
     def create_cluster(self, cluster, retry=DEFAULT, timeout=DEFAULT):
diff --git a/tests/contrib/hooks/test_gcp_container_hook.py b/tests/contrib/hooks/test_gcp_container_hook.py
index f3705ea4ce..6e13461395 100644
--- a/tests/contrib/hooks/test_gcp_container_hook.py
+++ b/tests/contrib/hooks/test_gcp_container_hook.py
@@ -61,6 +61,22 @@ def test_delete_cluster(self, wait_mock, convert_mock):
         wait_mock.assert_called_with(client_delete.return_value)
         convert_mock.assert_not_called()
 
+    @mock.patch(
+        "airflow.contrib.hooks.gcp_container_hook.GKEClusterHook.log")
+    @mock.patch("airflow.contrib.hooks.gcp_container_hook.GKEClusterHook._dict_to_proto")
+    @mock.patch(
+        "airflow.contrib.hooks.gcp_container_hook.GKEClusterHook.wait_for_operation")
+    def test_delete_cluster_not_found(self, wait_mock, convert_mock, log_mock):
+        from google.api_core.exceptions import NotFound
+        # To force an error
+        message = 'Not Found'
+        self.gke_hook.client.delete_cluster.side_effect = NotFound(message=message)
+
+        self.gke_hook.delete_cluster(None)
+        wait_mock.assert_not_called()
+        convert_mock.assert_not_called()
+        log_mock.info.assert_any_call("Assuming Success: " + message)
+
     @mock.patch("airflow.contrib.hooks.gcp_container_hook.GKEClusterHook._dict_to_proto")
     @mock.patch(
         "airflow.contrib.hooks.gcp_container_hook.GKEClusterHook.wait_for_operation")
@@ -107,7 +123,7 @@ def test_create_cluster_proto(self, wait_mock, convert_mock):
     @mock.patch("airflow.contrib.hooks.gcp_container_hook.GKEClusterHook._dict_to_proto")
     @mock.patch(
         "airflow.contrib.hooks.gcp_container_hook.GKEClusterHook.wait_for_operation")
-    def test_delete_cluster_dict(self, wait_mock, convert_mock):
+    def test_create_cluster_dict(self, wait_mock, convert_mock):
         mock_cluster_dict = {'name': CLUSTER_NAME}
         retry_mock, timeout_mock = mock.Mock(), mock.Mock()
 
@@ -135,6 +151,22 @@ def test_create_cluster_error(self, wait_mock, convert_mock):
             wait_mock.assert_not_called()
             convert_mock.assert_not_called()
 
+    @mock.patch(
+        "airflow.contrib.hooks.gcp_container_hook.GKEClusterHook.log")
+    @mock.patch("airflow.contrib.hooks.gcp_container_hook.GKEClusterHook._dict_to_proto")
+    @mock.patch(
+        "airflow.contrib.hooks.gcp_container_hook.GKEClusterHook.wait_for_operation")
+    def test_create_cluster_already_exists(self, wait_mock, convert_mock, log_mock):
+        from google.api_core.exceptions import AlreadyExists
+        # To force an error
+        message = 'Already Exists'
+        self.gke_hook.client.create_cluster.side_effect = AlreadyExists(message=message)
+
+        self.gke_hook.create_cluster({})
+        wait_mock.assert_not_called()
+        self.assertEquals(convert_mock.call_count, 1)
+        log_mock.info.assert_any_call("Assuming Success: " + message)
+
 
 class GKEClusterHookGetTest(unittest.TestCase):
     def setUp(self):


 

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
users@infra.apache.org


With regards,
Apache Git Services