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