def execute(self, context) -> str: hook = GKEHook( gcp_conn_id=self.gcp_conn_id, location=self.location, impersonation_chain=self.impersonation_chain, ) create_op = hook.create_cluster(cluster=self.body, project_id=self.project_id) return create_op
def execute(self, context: 'Context') -> str: hook = GKEHook( gcp_conn_id=self.gcp_conn_id, location=self.location, impersonation_chain=self.impersonation_chain, ) create_op = hook.create_cluster(cluster=self.body, project_id=self.project_id) KubernetesEngineClusterLink.persist(context=context, task_instance=self, cluster=self.body) return create_op
class TestGKEHookCreate(unittest.TestCase): def setUp(self): self.gke_hook = GKEHook(location=GKE_ZONE) self.gke_hook._client = mock.Mock() @mock.patch( 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', new_callable=PropertyMock, return_value=None, ) @mock.patch( "airflow.providers.google.cloud.hooks.kubernetes_engine.ParseDict") @mock.patch( "airflow.providers.google.cloud.hooks.kubernetes_engine.GKEHook.wait_for_operation" ) def test_create_cluster_proto(self, wait_mock, convert_mock, mock_project_id): 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(cluster=mock_cluster_proto, project_id=TEST_GCP_PROJECT_ID, retry=retry_mock, timeout=timeout_mock) client_create.assert_called_once_with( project_id=TEST_GCP_PROJECT_ID, zone=GKE_ZONE, cluster=mock_cluster_proto, retry=retry_mock, timeout=timeout_mock, ) wait_mock.assert_called_once_with(client_create.return_value) convert_mock.assert_not_called() @mock.patch( 'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id', new_callable=PropertyMock, return_value=None, ) @mock.patch( "airflow.providers.google.cloud.hooks.kubernetes_engine.ParseDict") @mock.patch( "airflow.providers.google.cloud.hooks.kubernetes_engine.GKEHook.wait_for_operation" ) def test_create_cluster_dict(self, wait_mock, convert_mock, mock_project_id): 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(cluster=mock_cluster_dict, project_id=TEST_GCP_PROJECT_ID, retry=retry_mock, timeout=timeout_mock) client_create.assert_called_once_with( project_id=TEST_GCP_PROJECT_ID, zone=GKE_ZONE, cluster=proto_mock, retry=retry_mock, timeout=timeout_mock, ) wait_mock.assert_called_once_with(client_create.return_value) convert_mock.assert_called_once_with({'name': 'test-cluster'}, Cluster()) @mock.patch( "airflow.providers.google.cloud.hooks.kubernetes_engine.ParseDict") @mock.patch( "airflow.providers.google.cloud.hooks.kubernetes_engine.GKEHook.wait_for_operation" ) def test_create_cluster_error(self, wait_mock, convert_mock): # to force an error mock_cluster_proto = None with pytest.raises(AirflowException): self.gke_hook.create_cluster(mock_cluster_proto) # pylint: disable=no-value-for-parameter wait_mock.assert_not_called() convert_mock.assert_not_called() @mock.patch( 'airflow.providers.google.cloud.hooks.kubernetes_engine.GKEHook._get_credentials_and_project_id', return_value=(mock.MagicMock(), TEST_GCP_PROJECT_ID), ) @mock.patch( "airflow.providers.google.cloud.hooks.kubernetes_engine.GKEHook.log") @mock.patch( "airflow.providers.google.cloud.hooks.kubernetes_engine.ParseDict") @mock.patch( "airflow.providers.google.cloud.hooks.kubernetes_engine.GKEHook.wait_for_operation" ) def test_create_cluster_already_exists(self, wait_mock, convert_mock, log_mock, mock_get_credentials): 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(cluster={}, project_id=TEST_GCP_PROJECT_ID) wait_mock.assert_not_called() assert convert_mock.call_count == 1 log_mock.info.assert_any_call("Assuming Success: %s", message)
def execute(self, context): hook = GKEHook(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