def map(message): """Map Step: Create COGs.""" if isinstance(message, str): message = json.loads(message) jobid = message["jobid"] partid = message["partid"] mosaicid = message["metadata"]["mosaicid"] logger.info(f"[MAP] Processing {jobid} - {partid}") with Part(jobid, partid, progress=progress): src_path = message["src_path"] # with rasterio.open(src_path) as src_dst: # witdh, height, dtype = src_dst.width, src_dst.height, src_dst.dtypes[0] # size = numpy.array([0], dtype=dtype).nbytes * witdh * height # logger.info(f"Size of image {size}") bname = os.path.splitext(os.path.basename(src_path))[0] out_key = os.path.join("cogs", mosaicid, f"{bname}_cog.tif") translator.process( src_path, os.environ["MOSAIC_BUCKET"], out_key, message["profile_name"], message["profile_options"], **message["options"], ) return True
def test_Part_job_not_done(aws_send_message, monkeypatch): monkeypatch.setenv('WorkTopic', 'abc123') monkeypatch.setenv('ProgressTable', 'arn::table/foo') with Part(jobid='123', partid=1, progress=MockProgress()): pass aws_send_message.assert_not_called()
def test_Part_job_done(aws_send_message, monkeypatch): monkeypatch.setenv('WorkTopic', 'abc123') monkeypatch.setenv('ProgressTable', 'arn::table/foo') class CustomProgress(MockProgress): def complete_part(self, jobid, partid): return True with Part(jobid='123', partid=1, progress=CustomProgress()): pass aws_send_message.assert_called_once() assert aws_send_message.call_args[1].get('subject') == 'reduce'
def test_Part_fail_job_on(status, fail_job, aws_send_message, monkeypatch): monkeypatch.setenv('WorkTopic', 'abc123') monkeypatch.setenv('ProgressTable', 'arn::table/foo') class CustomException(Exception): pass # the job WILL be marked as failed, the exeption is still re-raised with pytest.raises(CustomException): with Part(jobid=2, partid=2, fail_job_on=[CustomException]): raise CustomException() fail_job.assert_called_once_with(2, 2) aws_send_message.assert_not_called()
def test_progress_reduce_once(aws_send_message, sns_worker, monkeypatch): """Ensure that reduce message is only sent once, even if parts are completed multiple times """ monkeypatch.setenv('WorkTopic', 'abc123') progress = RedisProgress() jobid = create_job(parts, progress=progress) for i in range(3): with Part(jobid, i, progress=progress): pass aws_send_message.assert_called_once() assert aws_send_message.call_args[1].get('subject') == 'reduce' aws_send_message.reset_mock() # Finishing already completed parts should not trigger reduce message, only warn with pytest.warns(UserWarning) as record: for i in range(3): with Part(jobid, i, progress=progress): pass assert 'skip' in record[0].message.args[0] aws_send_message.assert_not_called()
def test_progress_all_done(aws_send_message, sns_worker, monkeypatch): monkeypatch.setenv('WorkTopic', 'abc123') progress = RedisProgress() jobid = create_job(parts, progress=progress) for i in range(3): with Part(jobid, i, progress=progress): pass assert progress.status(jobid)['remaining'] == 0 # 3 map messages and 1 reduce message sent sns_worker.assert_called_once() assert len(sns_worker.call_args[0][0]) == 3 assert sns_worker.call_args[1].get('subject') == 'map' aws_send_message.assert_called_once() assert aws_send_message.call_args[1].get('subject') == 'reduce'
def test_progress_one_done(aws_send_message, sns_worker, monkeypatch): monkeypatch.setenv('WorkTopic', 'abc123') sns_worker.side_effect = [True] progress = RedisProgress() jobid = create_job(parts, progress=progress) with Part(jobid, 0, progress=progress): pass assert progress.status(jobid)['remaining'] == 2 # 3 map messages, No reduce message sent sns_worker.assert_called_once() assert len(sns_worker.call_args[0][0]) == 3 assert sns_worker.call_args[1].get('subject') == 'map' aws_send_message.assert_not_called()
def test_Part_job_done_on_reduce_callback(monkeypatch): monkeypatch.setenv('WorkTopic', 'abc123') monkeypatch.setenv('ProgressTable', 'arn::table/foo') class CustomProgress(MockProgress): def complete_part(self, jobid, partid): return True on_reduce = Mock() with Part(jobid='123', partid=1, progress=CustomProgress(), on_reduce=on_reduce): pass on_reduce.assert_called_once() assert on_reduce.call_args[1].get('subject') == 'reduce'
def test_Part_already_failed(aws_send_message, monkeypatch): monkeypatch.setenv('WorkTopic', 'abc123') monkeypatch.setenv('ProgressTable', 'arn::table/foo') class CustomProgress(MockProgress): def status(self, x): return {'progress': 0.3, 'failed': True} class CustomException(Exception): pass msg = {'source': 'a.tif', 'partid': 1, 'jobid': '123'} with pytest.raises(JobFailed) as e: with Part(fail_job_on=[CustomException], progress=CustomProgress(), **msg): # sees that the overall job failed and will not execute this code raise NotImplementedError("You'll never get here") assert 'already failed' in str(e) aws_send_message.assert_not_called()
def doit(msg, jobid=jobid, progress=progress): partid = msg['partid'] with Part(partid=partid, jobid=jobid, progress=progress): pass
from watchbot_progress import Part, create_job parts = [{'source': 'a.tif'}, {'source': 'b.tif'}, {'source': 'c.tif'}] def annotate_parts(parts): return [ dict(**part, partid=partid, jobid=jobid) for partid, part in enumerate(parts) ] if __name__ == "__main__": # ------------- Subject: start-job jobid = create_job(parts) # ------------- Subject: map for msg in annotate_parts(parts): with Part(msg): partid = msg['partid'] print(f'PROCCESS {partid}') print('PROCCESS complete') print("complete")
def test_Part_bad_progress(aws_send_message): with pytest.raises(ProgressTypeError): with Part(partid=1, jobid='1', progress=Exception()): pass aws_send_message.assert_not_called()