def test_cancel_succeed(self, mock_dataproc_client, mock_kfp_context,
mock_dump_json, mock_display):
mock_kfp_context().__enter__().context_id.return_value = 'ctx1'
operation = {'name': 'mock_operation'}
mock_dataproc_client().create_cluster.return_value = (operation)
cluster = {
'projectId': 'mock_project',
'config': {},
'clusterName': 'my-cluster-ctx1'
}
mock_dataproc_client().wait_for_operation_done.return_value = ({
'response':
cluster
})
create_cluster('mock_project', 'mock-region')
cancel_func = mock_kfp_context.call_args[1]['on_cancel']
cancel_func()
mock_dataproc_client().cancel_operation.assert_called_with(
'mock_operation')
def test_create_cluster_succeed(self, mock_dataproc_client,
mock_kfp_context, mock_dump_json,
mock_display):
mock_kfp_context().__enter__().context_id.return_value = 'ctx1'
cluster = {
'projectId': 'mock_project',
'config': {},
'clusterName': 'cluster-ctx1'
}
mock_dataproc_client().wait_for_operation_done.return_value = ({
'response':
cluster
})
result = create_cluster('mock_project', 'mock-region')
self.assertEqual(cluster, result)
mock_dataproc_client().create_cluster.assert_called_with(
'mock_project', 'mock-region', cluster, request_id='ctx1')
def test_create_cluster_with_specs_succeed(self, mock_dataproc_client,
mock_kfp_context,
mock_dump_json, mock_display):
mock_kfp_context().__enter__().context_id.return_value = 'ctx1'
cluster = {
'projectId': 'mock_project',
'config': {
'initializationActions': [{
'executableFile': 'gs://action/1'
}, {
'executableFile': 'gs://action/2'
}],
'configBucket':
'gs://config/bucket'
},
'softwareConfig': {
'imageVersion': '1.10'
},
'labels': {
'label-1': 'value-1'
},
'clusterName': 'test-cluster'
}
mock_dataproc_client().wait_for_operation_done.return_value = ({
'response':
cluster
})
result = create_cluster(
'mock_project',
'mock-region',
name='test-cluster',
initialization_actions=['gs://action/1', 'gs://action/2'],
config_bucket='gs://config/bucket',
image_version='1.10',
cluster={'labels': {
'label-1': 'value-1'
}})
self.assertEqual(cluster, result)
mock_dataproc_client().create_cluster.assert_called_with(
'mock_project', 'mock-region', cluster, request_id='ctx1')