def execute(self, context): self._check_input() hook = GKEClusterHook(gcp_conn_id=self.gcp_conn_id, location=self.location) create_op = hook.create_cluster(cluster=self.body, project_id=self.project_id) return create_op
def execute(self, context): self._check_input() hook = GKEClusterHook(gcp_conn_id=self.gcp_conn_id, location=self.location) delete_result = hook.delete_cluster(name=self.name, project_id=self.project_id) return delete_result
def setUp(self): with mock.patch.object(GKEClusterHook, "__init__", return_value=None): self.gke_hook = GKEClusterHook(None, None, None) self.gke_hook.project_id = TEST_PROJECT_ID self.gke_hook.location = ZONE self.gke_hook.client = mock.Mock()
class GKEClusterHookDeleteTest(unittest.TestCase): def setUp(self): with mock.patch.object(GKEClusterHook, "__init__", return_value=None): self.gke_hook = GKEClusterHook(None, None, None) self.gke_hook.project_id = TEST_PROJECT_ID self.gke_hook.location = ZONE self.gke_hook.client = mock.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(self, wait_mock, convert_mock): retry_mock, timeout_mock = mock.Mock(), mock.Mock() client_delete = self.gke_hook.client.delete_cluster = mock.Mock() self.gke_hook.delete_cluster(name=CLUSTER_NAME, retry=retry_mock, timeout=timeout_mock) client_delete.assert_called_with(project_id=TEST_PROJECT_ID, zone=ZONE, cluster_id=CLUSTER_NAME, retry=retry_mock, timeout=timeout_mock) wait_mock.assert_called_with(client_delete.return_value) convert_mock.assert_not_called() @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_error(self, wait_mock, convert_mock): # To force an error self.gke_hook.client.delete_cluster.side_effect = AirflowException( '400') with self.assertRaises(AirflowException): self.gke_hook.delete_cluster(None) wait_mock.assert_not_called() convert_mock.assert_not_called()
class GKEClusterHookGetTest(unittest.TestCase): def setUp(self): self.gke_hook = GKEClusterHook(location=ZONE) self.gke_hook._client = mock.Mock() def test_get_cluster(self): retry_mock, timeout_mock = mock.Mock(), mock.Mock() client_get = self.gke_hook._client.get_cluster = mock.Mock() self.gke_hook.get_cluster(name=CLUSTER_NAME, project_id=TEST_PROJECT_ID, retry=retry_mock, timeout=timeout_mock) client_get.assert_called_with(project_id=TEST_PROJECT_ID, zone=ZONE, cluster_id=CLUSTER_NAME, retry=retry_mock, timeout=timeout_mock)
class GKEClusterHookGetTest(unittest.TestCase): def setUp(self): self.gke_hook = GKEClusterHook(location=GKE_ZONE) self.gke_hook._client = mock.Mock() def test_get_cluster(self): retry_mock, timeout_mock = mock.Mock(), mock.Mock() client_get = self.gke_hook._client.get_cluster = mock.Mock() self.gke_hook.get_cluster(name=CLUSTER_NAME, project_id=TEST_GCP_PROJECT_ID, retry=retry_mock, timeout=timeout_mock) client_get.assert_called_with(project_id=TEST_GCP_PROJECT_ID, zone=GKE_ZONE, cluster_id=CLUSTER_NAME, retry=retry_mock, timeout=timeout_mock)
class GKEClusterHookGetTest(unittest.TestCase): def setUp(self): with mock.patch.object(GKEClusterHook, "__init__", return_value=None): self.gke_hook = GKEClusterHook(None, None, None) self.gke_hook.project_id = TEST_PROJECT_ID self.gke_hook.location = ZONE self.gke_hook.client = mock.Mock() def test_get_cluster(self): retry_mock, timeout_mock = mock.Mock(), mock.Mock() client_get = self.gke_hook.client.get_cluster = mock.Mock() self.gke_hook.get_cluster(name=CLUSTER_NAME, retry=retry_mock, timeout=timeout_mock) client_get.assert_called_with(project_id=TEST_PROJECT_ID, zone=ZONE, cluster_id=CLUSTER_NAME, retry=retry_mock, timeout=timeout_mock)
class GKEClusterHookCreateTest(unittest.TestCase): def setUp(self): with mock.patch.object(GKEClusterHook, "__init__", return_value=None): self.gke_hook = GKEClusterHook(None, None, None) self.gke_hook.project_id = TEST_PROJECT_ID self.gke_hook.location = ZONE self.gke_hook.client = mock.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_create_cluster_proto(self, wait_mock, convert_mock): from google.cloud.container_v1.proto.cluster_service_pb2 import Cluster mock_cluster_proto = Cluster() mock_cluster_proto.name = CLUSTER_NAME retry_mock, timeout_mock = mock.Mock(), mock.Mock() client_create = self.gke_hook.client.create_cluster = mock.Mock() self.gke_hook.create_cluster(mock_cluster_proto, retry=retry_mock, timeout=timeout_mock) client_create.assert_called_with(project_id=TEST_PROJECT_ID, zone=ZONE, cluster=mock_cluster_proto, retry=retry_mock, timeout=timeout_mock) wait_mock.assert_called_with(client_create.return_value) convert_mock.assert_not_called() @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): mock_cluster_dict = {'name': CLUSTER_NAME} retry_mock, timeout_mock = mock.Mock(), mock.Mock() client_create = self.gke_hook.client.create_cluster = mock.Mock() proto_mock = convert_mock.return_value = mock.Mock() self.gke_hook.create_cluster(mock_cluster_dict, retry=retry_mock, timeout=timeout_mock) client_create.assert_called_with(project_id=TEST_PROJECT_ID, zone=ZONE, cluster=proto_mock, retry=retry_mock, timeout=timeout_mock) wait_mock.assert_called_with(client_create.return_value) self.assertEqual(convert_mock.call_args[1]['py_dict'], mock_cluster_dict) @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_error(self, wait_mock, convert_mock): # to force an error mock_cluster_proto = None with self.assertRaises(AirflowException): self.gke_hook.create_cluster(mock_cluster_proto) wait_mock.assert_not_called() convert_mock.assert_not_called()
class GKEClusterHookGetTest(unittest.TestCase): def setUp(self): with mock.patch.object(GKEClusterHook, "__init__", return_value=None): self.gke_hook = GKEClusterHook(None, None, None) self.gke_hook.project_id = TEST_PROJECT_ID self.gke_hook.location = ZONE self.gke_hook.client = mock.Mock() def test_get_cluster(self): retry_mock, timeout_mock = mock.Mock(), mock.Mock() client_get = self.gke_hook.client.get_cluster = mock.Mock() self.gke_hook.get_cluster(name=CLUSTER_NAME, retry=retry_mock, timeout=timeout_mock) client_get.assert_called_with(project_id=TEST_PROJECT_ID, zone=ZONE, cluster_id=CLUSTER_NAME, retry=retry_mock, timeout=timeout_mock)
class GKEClusterHookDeleteTest(unittest.TestCase): def setUp(self): with mock.patch.object(GKEClusterHook, "__init__", return_value=None): self.gke_hook = GKEClusterHook(None, None, None) self.gke_hook.project_id = TEST_PROJECT_ID self.gke_hook.location = ZONE self.gke_hook.client = mock.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(self, wait_mock, convert_mock): retry_mock, timeout_mock = mock.Mock(), mock.Mock() client_delete = self.gke_hook.client.delete_cluster = mock.Mock() self.gke_hook.delete_cluster(name=CLUSTER_NAME, retry=retry_mock, timeout=timeout_mock) client_delete.assert_called_with(project_id=TEST_PROJECT_ID, zone=ZONE, cluster_id=CLUSTER_NAME, retry=retry_mock, timeout=timeout_mock) wait_mock.assert_called_with(client_delete.return_value) convert_mock.assert_not_called() @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_error(self, wait_mock, convert_mock): # To force an error self.gke_hook.client.delete_cluster.side_effect = AirflowException('400') with self.assertRaises(AirflowException): self.gke_hook.delete_cluster(None) wait_mock.assert_not_called() convert_mock.assert_not_called()
class GKEClusterHookDeleteTest(unittest.TestCase): def setUp(self): self.gke_hook = GKEClusterHook(location=GKE_ZONE) self.gke_hook._client = mock.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(self, wait_mock, convert_mock): retry_mock, timeout_mock = mock.Mock(), mock.Mock() client_delete = self.gke_hook._client.delete_cluster = mock.Mock() self.gke_hook.delete_cluster(name=CLUSTER_NAME, project_id=TEST_GCP_PROJECT_ID, retry=retry_mock, timeout=timeout_mock) client_delete.assert_called_with(project_id=TEST_GCP_PROJECT_ID, zone=GKE_ZONE, cluster_id=CLUSTER_NAME, retry=retry_mock, timeout=timeout_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('not-existing') wait_mock.assert_not_called() convert_mock.assert_not_called() log_mock.info.assert_any_call("Assuming Success: %s", 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") def test_delete_cluster_error(self, wait_mock, convert_mock): # To force an error self.gke_hook._client.delete_cluster.side_effect = AirflowException('400') with self.assertRaises(AirflowException): self.gke_hook.delete_cluster('a-cluster') wait_mock.assert_not_called() convert_mock.assert_not_called()
class GKEClusterHookDeleteTest(unittest.TestCase): def setUp(self): self.gke_hook = GKEClusterHook(location=ZONE) self.gke_hook._client = mock.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(self, wait_mock, convert_mock): retry_mock, timeout_mock = mock.Mock(), mock.Mock() client_delete = self.gke_hook._client.delete_cluster = mock.Mock() self.gke_hook.delete_cluster(name=CLUSTER_NAME, project_id=TEST_PROJECT_ID, retry=retry_mock, timeout=timeout_mock) client_delete.assert_called_with(project_id=TEST_PROJECT_ID, zone=ZONE, cluster_id=CLUSTER_NAME, retry=retry_mock, timeout=timeout_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") def test_delete_cluster_error(self, wait_mock, convert_mock): # To force an error self.gke_hook._client.delete_cluster.side_effect = AirflowException('400') with self.assertRaises(AirflowException): self.gke_hook.delete_cluster(None) wait_mock.assert_not_called() convert_mock.assert_not_called()
def setUp(self): self.gke_hook = GKEClusterHook(location=GKE_ZONE) self.gke_hook._client = mock.Mock()
class GKEClusterHookCreateTest(unittest.TestCase): def setUp(self): self.gke_hook = GKEClusterHook(location=GKE_ZONE) self.gke_hook._client = mock.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_create_cluster_proto(self, wait_mock, convert_mock): from google.cloud.container_v1.proto.cluster_service_pb2 import Cluster mock_cluster_proto = Cluster() mock_cluster_proto.name = CLUSTER_NAME retry_mock, timeout_mock = mock.Mock(), mock.Mock() client_create = self.gke_hook._client.create_cluster = mock.Mock() self.gke_hook.create_cluster(mock_cluster_proto, project_id=TEST_GCP_PROJECT_ID, retry=retry_mock, timeout=timeout_mock) client_create.assert_called_with(project_id=TEST_GCP_PROJECT_ID, zone=GKE_ZONE, cluster=mock_cluster_proto, retry=retry_mock, timeout=timeout_mock) wait_mock.assert_called_with(client_create.return_value) convert_mock.assert_not_called() @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_dict(self, wait_mock, convert_mock): mock_cluster_dict = {'name': CLUSTER_NAME} retry_mock, timeout_mock = mock.Mock(), mock.Mock() client_create = self.gke_hook._client.create_cluster = mock.Mock() proto_mock = convert_mock.return_value = mock.Mock() self.gke_hook.create_cluster(mock_cluster_dict, project_id=TEST_GCP_PROJECT_ID, retry=retry_mock, timeout=timeout_mock) client_create.assert_called_with(project_id=TEST_GCP_PROJECT_ID, zone=GKE_ZONE, cluster=proto_mock, retry=retry_mock, timeout=timeout_mock) wait_mock.assert_called_with(client_create.return_value) self.assertEqual(convert_mock.call_args[1]['py_dict'], mock_cluster_dict) @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_error(self, wait_mock, convert_mock): # to force an error mock_cluster_proto = None with self.assertRaises(AirflowException): self.gke_hook.create_cluster(mock_cluster_proto) 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.assertEqual(convert_mock.call_count, 1) log_mock.info.assert_any_call("Assuming Success: %s", message)
class GKEClusterHookTest(unittest.TestCase): def setUp(self): self.gke_hook = GKEClusterHook(location=GKE_ZONE) self.gke_hook._client = mock.Mock() @mock.patch('airflow.contrib.hooks.gcp_container_hook.container_v1.' 'ClusterManagerClient') @mock.patch('airflow.contrib.hooks.gcp_container_hook.ClientInfo') @mock.patch('airflow.contrib.hooks.gcp_container_hook.GKEClusterHook._get_credentials') def test_get_client(self, mock_get_credentials, mock_client_info, mock_client): self.gke_hook._client = None self.gke_hook.get_client() mock_get_credentials.assert_called() mock_client.assert_called_with( credentials=mock_get_credentials.return_value, client_info=mock_client_info.return_value) def test_get_operation(self): self.gke_hook._client.get_operation = mock.Mock() self.gke_hook.get_operation('TEST_OP', project_id=TEST_GCP_PROJECT_ID) self.gke_hook._client.get_operation.assert_called_with( project_id=TEST_GCP_PROJECT_ID, zone=GKE_ZONE, operation_id='TEST_OP') def test_append_label(self): key = 'test-key' val = 'test-val' mock_proto = mock.Mock() self.gke_hook._append_label(mock_proto, key, val) mock_proto.resource_labels.update.assert_called_with({key: val}) def test_append_label_replace(self): key = 'test-key' val = 'test.val+this' mock_proto = mock.Mock() self.gke_hook._append_label(mock_proto, key, val) mock_proto.resource_labels.update.assert_called_with({key: 'test-val-this'}) @mock.patch("time.sleep") def test_wait_for_response_done(self, time_mock): from google.cloud.container_v1.gapic.enums import Operation mock_op = mock.Mock() mock_op.status = Operation.Status.DONE self.gke_hook.wait_for_operation(mock_op) self.assertEqual(time_mock.call_count, 1) @mock.patch("time.sleep") def test_wait_for_response_exception(self, time_mock): from google.cloud.container_v1.gapic.enums import Operation from google.cloud.exceptions import GoogleCloudError mock_op = mock.Mock() mock_op.status = Operation.Status.ABORTING with self.assertRaises(GoogleCloudError): self.gke_hook.wait_for_operation(mock_op) self.assertEqual(time_mock.call_count, 1) @mock.patch("airflow.contrib.hooks.gcp_container_hook.GKEClusterHook.get_operation") @mock.patch("time.sleep") def test_wait_for_response_running(self, time_mock, operation_mock): from google.cloud.container_v1.gapic.enums import Operation running_op, done_op, pending_op = mock.Mock(), mock.Mock(), mock.Mock() running_op.status = Operation.Status.RUNNING done_op.status = Operation.Status.DONE pending_op.status = Operation.Status.PENDING # Status goes from Running -> Pending -> Done operation_mock.side_effect = [pending_op, done_op] self.gke_hook.wait_for_operation(running_op, project_id=TEST_GCP_PROJECT_ID) self.assertEqual(time_mock.call_count, 3) operation_mock.assert_any_call(running_op.name, project_id=TEST_GCP_PROJECT_ID) operation_mock.assert_any_call(pending_op.name, project_id=TEST_GCP_PROJECT_ID) self.assertEqual(operation_mock.call_count, 2) @mock.patch("google.protobuf.json_format.Parse") @mock.patch("json.dumps") def test_dict_to_proto(self, dumps_mock, parse_mock): mock_dict = {'name': 'test'} mock_proto = mock.Mock() dumps_mock.return_value = mock.Mock() self.gke_hook._dict_to_proto(mock_dict, mock_proto) dumps_mock.assert_called_with(mock_dict) parse_mock.assert_called_with(dumps_mock(), mock_proto)
def execute(self, context): self._check_input() hook = GKEClusterHook(self.project_id, self.location) delete_result = hook.delete_cluster(name=self.name) return delete_result
def execute(self, context): self._check_input() hook = GKEClusterHook(self.project_id, self.location) create_op = hook.create_cluster(cluster=self.body) return create_op
def setUp(self): with mock.patch.object(GKEClusterHook, "__init__", return_value=None): self.gke_hook = GKEClusterHook(None, None, None) self.gke_hook.project_id = TEST_PROJECT_ID self.gke_hook.location = ZONE self.gke_hook.client = mock.Mock()
class GKEClusterHookTest(unittest.TestCase): def setUp(self): self.gke_hook = GKEClusterHook(location=ZONE) self.gke_hook._client = mock.Mock() @mock.patch('airflow.contrib.hooks.gcp_container_hook.container_v1.ClusterManagerClient') @mock.patch('airflow.contrib.hooks.gcp_container_hook.ClientInfo') @mock.patch('airflow.contrib.hooks.gcp_container_hook.GKEClusterHook._get_credentials') def test_get_client(self, mock_get_credentials, mock_client_info, mock_client): self.gke_hook._client = None self.gke_hook.get_client() mock_get_credentials.assert_called() mock_client.assert_called_with( credentials=mock_get_credentials.return_value, client_info=mock_client_info.return_value) def test_get_operation(self): self.gke_hook._client.get_operation = mock.Mock() self.gke_hook.get_operation('TEST_OP', project_id=TEST_PROJECT_ID) self.gke_hook._client.get_operation.assert_called_with(project_id=TEST_PROJECT_ID, zone=ZONE, operation_id='TEST_OP') def test_append_label(self): key = 'test-key' val = 'test-val' mock_proto = mock.Mock() self.gke_hook._append_label(mock_proto, key, val) mock_proto.resource_labels.update.assert_called_with({key: val}) def test_append_label_replace(self): key = 'test-key' val = 'test.val+this' mock_proto = mock.Mock() self.gke_hook._append_label(mock_proto, key, val) mock_proto.resource_labels.update.assert_called_with({key: 'test-val-this'}) @mock.patch("time.sleep") def test_wait_for_response_done(self, time_mock): from google.cloud.container_v1.gapic.enums import Operation mock_op = mock.Mock() mock_op.status = Operation.Status.DONE self.gke_hook.wait_for_operation(mock_op) self.assertEqual(time_mock.call_count, 1) @mock.patch("time.sleep") def test_wait_for_response_exception(self, time_mock): from google.cloud.container_v1.gapic.enums import Operation from google.cloud.exceptions import GoogleCloudError mock_op = mock.Mock() mock_op.status = Operation.Status.ABORTING with self.assertRaises(GoogleCloudError): self.gke_hook.wait_for_operation(mock_op) self.assertEqual(time_mock.call_count, 1) @mock.patch("airflow.contrib.hooks.gcp_container_hook.GKEClusterHook.get_operation") @mock.patch("time.sleep") def test_wait_for_response_running(self, time_mock, operation_mock): from google.cloud.container_v1.gapic.enums import Operation running_op, done_op, pending_op = mock.Mock(), mock.Mock(), mock.Mock() running_op.status = Operation.Status.RUNNING done_op.status = Operation.Status.DONE pending_op.status = Operation.Status.PENDING # Status goes from Running -> Pending -> Done operation_mock.side_effect = [pending_op, done_op] self.gke_hook.wait_for_operation(running_op, project_id=TEST_PROJECT_ID) self.assertEqual(time_mock.call_count, 3) operation_mock.assert_any_call(running_op.name, project_id=TEST_PROJECT_ID) operation_mock.assert_any_call(pending_op.name, project_id=TEST_PROJECT_ID) self.assertEqual(operation_mock.call_count, 2) @mock.patch("google.protobuf.json_format.Parse") @mock.patch("json.dumps") def test_dict_to_proto(self, dumps_mock, parse_mock): mock_dict = {'name': 'test'} mock_proto = mock.Mock() dumps_mock.return_value = mock.Mock() self.gke_hook._dict_to_proto(mock_dict, mock_proto) dumps_mock.assert_called_with(mock_dict) parse_mock.assert_called_with(dumps_mock(), mock_proto)
class GKEClusterHookCreateTest(unittest.TestCase): def setUp(self): with mock.patch.object(GKEClusterHook, "__init__", return_value=None): self.gke_hook = GKEClusterHook(None, None, None) self.gke_hook.project_id = TEST_PROJECT_ID self.gke_hook.location = ZONE self.gke_hook.client = mock.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_create_cluster_proto(self, wait_mock, convert_mock): from google.cloud.container_v1.proto.cluster_service_pb2 import Cluster mock_cluster_proto = Cluster() mock_cluster_proto.name = CLUSTER_NAME retry_mock, timeout_mock = mock.Mock(), mock.Mock() client_create = self.gke_hook.client.create_cluster = mock.Mock() self.gke_hook.create_cluster(mock_cluster_proto, retry=retry_mock, timeout=timeout_mock) client_create.assert_called_with(project_id=TEST_PROJECT_ID, zone=ZONE, cluster=mock_cluster_proto, retry=retry_mock, timeout=timeout_mock) wait_mock.assert_called_with(client_create.return_value) convert_mock.assert_not_called() @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): mock_cluster_dict = {'name': CLUSTER_NAME} retry_mock, timeout_mock = mock.Mock(), mock.Mock() client_create = self.gke_hook.client.create_cluster = mock.Mock() proto_mock = convert_mock.return_value = mock.Mock() self.gke_hook.create_cluster(mock_cluster_dict, retry=retry_mock, timeout=timeout_mock) client_create.assert_called_with(project_id=TEST_PROJECT_ID, zone=ZONE, cluster=proto_mock, retry=retry_mock, timeout=timeout_mock) wait_mock.assert_called_with(client_create.return_value) self.assertEqual(convert_mock.call_args[1]['py_dict'], mock_cluster_dict) @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_error(self, wait_mock, convert_mock): # to force an error mock_cluster_proto = None with self.assertRaises(AirflowException): self.gke_hook.create_cluster(mock_cluster_proto) wait_mock.assert_not_called() convert_mock.assert_not_called()
def setUp(self): self.gke_hook = GKEClusterHook(location=ZONE) self.gke_hook._client = mock.Mock()
class GKEClusterHookCreateTest(unittest.TestCase): def setUp(self): self.gke_hook = GKEClusterHook(location=ZONE) self.gke_hook._client = mock.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_create_cluster_proto(self, wait_mock, convert_mock): from google.cloud.container_v1.proto.cluster_service_pb2 import Cluster mock_cluster_proto = Cluster() mock_cluster_proto.name = CLUSTER_NAME retry_mock, timeout_mock = mock.Mock(), mock.Mock() client_create = self.gke_hook._client.create_cluster = mock.Mock() self.gke_hook.create_cluster(mock_cluster_proto, project_id=TEST_PROJECT_ID, retry=retry_mock, timeout=timeout_mock) client_create.assert_called_with(project_id=TEST_PROJECT_ID, zone=ZONE, cluster=mock_cluster_proto, retry=retry_mock, timeout=timeout_mock) wait_mock.assert_called_with(client_create.return_value) convert_mock.assert_not_called() @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_dict(self, wait_mock, convert_mock): mock_cluster_dict = {'name': CLUSTER_NAME} retry_mock, timeout_mock = mock.Mock(), mock.Mock() client_create = self.gke_hook._client.create_cluster = mock.Mock() proto_mock = convert_mock.return_value = mock.Mock() self.gke_hook.create_cluster(mock_cluster_dict, project_id=TEST_PROJECT_ID, retry=retry_mock, timeout=timeout_mock) client_create.assert_called_with(project_id=TEST_PROJECT_ID, zone=ZONE, cluster=proto_mock, retry=retry_mock, timeout=timeout_mock) wait_mock.assert_called_with(client_create.return_value) self.assertEqual(convert_mock.call_args[1]['py_dict'], mock_cluster_dict) @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_error(self, wait_mock, convert_mock): # to force an error mock_cluster_proto = None with self.assertRaises(AirflowException): self.gke_hook.create_cluster(mock_cluster_proto) 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)