[AIRFLOW-2863] Fix GKEClusterHook catching wrong exception (apache#3711)

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):

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):