def test_dataproc_batch_remote_runner_existing_operation_not_found( self, mock_get_requests, _, mock_auth_default): with open(self._gcp_resources, 'w') as f: f.write( f'{{"resources": [{{"resourceType": "DataprocLro", "resourceUri": "https://dataproc.googleapis.com/v1/projects/{self._project}/regions/{self._location}/operations/{self._operation_id}"}}]}}' ) mock_creds = mock.Mock() mock_creds.token = self._creds_token mock_auth_default.return_value = [mock_creds, 'project'] mock_polled_lro = mock.Mock(spec=requests.models.Response) mock_polled_lro.raise_for_status.side_effect = requests.exceptions.HTTPError( mock.Mock(status=404), 'Not found') mock_get_requests.return_value = mock_polled_lro with self.assertRaises(RuntimeError): dataproc_batch_remote_runner.create_spark_batch( type='SparkBatch', project=self._project, location=self._location, batch_id=self._batch_id, payload=json.dumps(self._test_spark_batch), gcp_resources=self._gcp_resources) mock_get_requests.assert_called_once_with( self=mock.ANY, url= f'https://dataproc.googleapis.com/v1/projects/{self._project}/regions/{self._location}/operations/{self._operation_id}' )
def test_dataproc_batch_remote_runner_spark_batch_succeeded( self, mock_time_sleep, mock_post_requests, mock_get_requests, _, mock_auth_default): mock_creds = mock.Mock(spec=google.auth.credentials.Credentials) mock_creds.token = self._creds_token mock_auth_default.return_value = [mock_creds, 'project'] mock_operation = mock.Mock(spec=requests.models.Response) mock_operation.json.return_value = { 'name': f'projects/{self._project}/regions/{self._location}/operations/{self._operation_id}', 'metadata': { 'batch': f'projects/{self._project}/locations/{self._location}/batches/{self._batch_id}' } } mock_post_requests.return_value = mock_operation mock_polled_lro = mock.Mock(spec=requests.models.Response) mock_polled_lro.json.return_value = { 'name': f'projects/{self._project}/regions/{self._location}/operations/{self._operation_id}', 'done': True, 'response': { 'name': f'projects/{self._project}/locations/{self._location}/batches/{self._batch_id}', 'state': 'SUCCEEDED' } } mock_get_requests.return_value = mock_polled_lro dataproc_batch_remote_runner.create_spark_batch( type='SparkBatch', project=self._project, location=self._location, batch_id=self._batch_id, payload=json.dumps(self._test_spark_batch), gcp_resources=self._gcp_resources) mock_post_requests.assert_called_once_with( self=mock.ANY, url= f'{self._dataproc_uri_prefix}/projects/{self._project}/locations/{self._location}/batches/?batchId={self._batch_id}', data=json.dumps(self._test_spark_batch), headers={'Authorization': f'Bearer {self._creds_token}'}) mock_get_requests.assert_called_once_with( self=mock.ANY, url= f'{self._dataproc_uri_prefix}/projects/{self._project}/regions/{self._location}/operations/{self._operation_id}', headers={'Authorization': f'Bearer {self._creds_token}'}) mock_time_sleep.assert_called_once() self._validate_gcp_resources_succeeded()
def test_dataproc_batch_remote_runner_spark_batch_exception_on_error( self, mock_time_sleep, mock_post_requests, mock_get_requests, _, mock_auth_default): mock_creds = mock.Mock() mock_creds.token = self._creds_token mock_auth_default.return_value = [mock_creds, 'project'] mock_operation = mock.Mock(spec=requests.models.Response) mock_operation.json.return_value = { 'name': f'projects/{self._project}/regions/{self._location}/operations/{self._operation_id}', 'metadata': { 'batch': f'projects/{self._project}/locations/{self._location}/batches/{self._batch_id}' } } mock_post_requests.return_value = mock_operation mock_polled_lro = mock.Mock() mock_polled_lro.json.return_value = { 'name': f'projects/{self._project}/regions/{self._location}/operations/{self._operation_id}', 'done': True, 'error': { 'code': 10 } } mock_get_requests.return_value = mock_polled_lro with self.assertRaises(RuntimeError): dataproc_batch_remote_runner.create_spark_batch( type='SparkBatch', project=self._project, location=self._location, batch_id=self._batch_id, payload=json.dumps(self._test_spark_batch), gcp_resources=self._gcp_resources) mock_post_requests.assert_called_once_with( self=mock.ANY, url= f'https://dataproc.googleapis.com/v1/projects/{self._project}/locations/{self._location}/batches/?batchId={self._batch_id}', data=json.dumps(self._test_spark_batch)) mock_get_requests.assert_called_once_with( self=mock.ANY, url= f'https://dataproc.googleapis.com/v1/projects/{self._project}/regions/{self._location}/operations/{self._operation_id}' ) mock_time_sleep.assert_called_once() self._validate_gcp_resources_succeeded()
def test_dataproc_batch_remote_runner_operation_exists_wrong_format( self, mock_auth_default): with open(self._gcp_resources, 'w') as f: f.write( '{"resources": [{"resourceType": "DataprocLro", "resourceUri": "https://dataproc.googleapis.com/v1/projects/test-project/regions/test-location/operations"}]}' ) mock_creds = mock.Mock() mock_creds.token = self._creds_token mock_auth_default.return_value = [mock_creds, 'project'] with self.assertRaises(ValueError): dataproc_batch_remote_runner.create_spark_batch( type='SparkBatch', project=self._project, location=self._location, batch_id=self._batch_id, payload=json.dumps(self._test_spark_batch), gcp_resources=self._gcp_resources)
def test_dataproc_batch_remote_runner_exception_with_more_than_one_resource_in_gcp_resources( self, mock_auth_default): with open(self._gcp_resources, 'w') as f: f.write(( f'{{"resources": [' f'{{"resourceType": "DataprocLro", "resourceUri": "{self._dataproc_uri_prefix}/projects/{self._project}/regions/{self._location}/operations/{self._operation_id}"}},' f'{{"resourceType": "DataprocLro", "resourceUri": "{self._dataproc_uri_prefix}/projects/{self._project}/regions/{self._location}/operations/{self._operation_id}"}}' f']}}')) mock_creds = mock.Mock(spec=google.auth.credentials.Credentials) mock_creds.token = self._creds_token mock_auth_default.return_value = [mock_creds, 'project'] with self.assertRaises(ValueError): dataproc_batch_remote_runner.create_spark_batch( type='SparkBatch', project=self._project, location=self._location, batch_id=self._batch_id, payload=json.dumps(self._test_spark_batch), gcp_resources=self._gcp_resources)
def test_dataproc_batch_remote_runner_existing_operation_succeeded( self, mock_get_requests, _, mock_auth_default): with open(self._gcp_resources, 'w') as f: f.write( f'{{"resources": [{{"resourceType": "DataprocLro", "resourceUri": "https://dataproc.googleapis.com/v1/projects/{self._project}/regions/{self._location}/operations/{self._operation_id}"}}]}}' ) mock_creds = mock.Mock() mock_creds.token = self._creds_token mock_auth_default.return_value = [mock_creds, 'project'] mock_polled_lro = mock.Mock() mock_polled_lro.json.return_value = { 'name': f'projects/{self._project}/regions/{self._location}/operations/{self._operation_id}', 'done': True, 'response': { 'name': f'projects/{self._project}/locations/{self._location}/batches/{self._batch_id}', 'state': 'SUCCEEDED' } } mock_get_requests.return_value = mock_polled_lro dataproc_batch_remote_runner.create_spark_batch( type='SparkBatch', project=self._project, location=self._location, batch_id=self._batch_id, payload=json.dumps(self._test_spark_batch), gcp_resources=self._gcp_resources) mock_get_requests.assert_called_once_with( self=mock.ANY, url= f'https://dataproc.googleapis.com/v1/projects/{self._project}/regions/{self._location}/operations/{self._operation_id}' )
def test_dataproc_batch_remote_runner_existing_operation_failed( self, mock_get_requests, _, mock_auth_default): with open(self._gcp_resources, 'w') as f: f.write( f'{{"resources": [{{"resourceType": "DataprocLro", "resourceUri": "{self._dataproc_uri_prefix}/projects/{self._project}/regions/{self._location}/operations/{self._operation_id}"}}]}}' ) mock_creds = mock.Mock(spec=google.auth.credentials.Credentials) mock_creds.token = self._creds_token mock_auth_default.return_value = [mock_creds, 'project'] mock_polled_lro = mock.Mock(spec=requests.models.Response) mock_polled_lro.json.return_value = { 'name': f'projects/{self._project}/regions/{self._location}/operations/{self._operation_id}', 'done': True, 'error': { 'code': 10 } } mock_get_requests.return_value = mock_polled_lro with self.assertRaises(RuntimeError): dataproc_batch_remote_runner.create_spark_batch( type='SparkBatch', project=self._project, location=self._location, batch_id=self._batch_id, payload=json.dumps(self._test_spark_batch), gcp_resources=self._gcp_resources) mock_get_requests.assert_called_once_with( self=mock.ANY, url= f'{self._dataproc_uri_prefix}/projects/{self._project}/regions/{self._location}/operations/{self._operation_id}', headers={'Authorization': f'Bearer {self._creds_token}'})