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