def test_get_beam_run_result_with_result_batches(self) -> None:
beam_job_models.BeamJobRunResultModel(job_id='123', stdout='abc').put()
beam_job_models.BeamJobRunResultModel(job_id='123', stderr='123').put()
beam_job_models.BeamJobRunResultModel(
job_id='123', stdout='def', stderr='456').put()
beam_job_run_result = beam_job_services.get_beam_job_run_result('123')
self.assertItemsEqual( # type: ignore[no-untyped-call]
beam_job_run_result.stdout.split('\n'), ['abc', 'def'])
self.assertItemsEqual( # type: ignore[no-untyped-call]
beam_job_run_result.stderr.split('\n'), ['123', '456'])
def test_get_beam_run_result(self) -> None:
beam_job_models.BeamJobRunResultModel(
job_id='123', stdout='abc', stderr='def').put()
beam_job_run_result = beam_job_services.get_beam_job_run_result('123')
self.assertEqual(beam_job_run_result.stdout, 'abc')
self.assertEqual(beam_job_run_result.stderr, 'def')
def test_get_new_id_raises_error_after_too_many_failed_attempts(
self) -> None:
model = beam_job_models.BeamJobRunResultModel(
id=beam_job_models.BeamJobRunResultModel.get_new_id(),
job_id='123')
model.update_timestamps()
model.put()
collision_context = self.swap_to_always_return(utils,
'convert_to_hash',
value=model.id)
with collision_context:
self.assertRaisesRegex( # type: ignore[no-untyped-call]
RuntimeError,
r'Failed to generate a unique ID after \d+ attempts',
beam_job_models.BeamJobRunResultModel.get_new_id)
def create_beam_job_run_result_model(
job_id: str, stdout: str, stderr: str
) -> beam_job_models.BeamJobRunResultModel:
"""Creates a new BeamJobRunResultModel without putting it into storage.
Args:
job_id: str. The ID of the job run to fetch.
stdout: str. The standard output from a job run.
stderr: str. The error output from a job run.
Returns:
BeamJobRunResultModel. The model.
"""
model_id = beam_job_models.BeamJobRunResultModel.get_new_id()
model = beam_job_models.BeamJobRunResultModel(
id=model_id, job_id=job_id, stdout=stdout, stderr=stderr)
model.update_timestamps()
return model
