def main(backend, number):
if backend == 'redis':
progress = RedisProgress(host='localhost', port='6379')
elif backend == 'dynamodb':
progress = DynamoProgress()
parts = [dict() for _ in range(number)]
workers = 25
# Simulates Subject:start messages
jobid = create_job(parts, progress=progress, workers=workers)
# Simulates concurrent Subject:map messages
# The annonate_parts function recreates the messages that would be recieved
msgs = annotate_parts(parts, jobid)
def doit(msg, jobid=jobid, progress=progress):
partid = msg['partid']
with Part(partid=partid, jobid=jobid, progress=progress):
pass
with futures.ThreadPoolExecutor(max_workers=workers) as executor:
executor.map(doit, msgs)
print(f"completed watchbot progress for {number} parts using {progress}")
def test_create_jobs(sns_worker, monkeypatch):
monkeypatch.setenv('WorkTopic', 'abc123')
monkeypatch.setenv('ProgressTable', 'arn::table/foo')
sns_worker.side_effect = [True]
with patch('uuid.uuid4', new=lambda: '000-123'):
assert create_job(parts, progress=MockProgress()) == '000-123'
sns_worker.assert_called_once()
assert len(sns_worker.call_args[0][0]) == 3
assert sns_worker.call_args[1].get('subject') == 'map'
def test_create_jobs_metadata(sns_worker, monkeypatch):
monkeypatch.setenv('WorkTopic', 'abc123')
monkeypatch.setenv('ProgressTable', 'arn::table/foo')
meta = {'foo': 'bar'}
with patch('uuid.uuid4', new=lambda: '000-123'):
assert create_job(parts, progress=MockProgress(),
metadata=meta) == '000-123'
sns_worker.assert_called_once()
assert len(sns_worker.call_args[0][0]) == 3
assert sns_worker.call_args[1].get('subject') == 'map'
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_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()
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_create_bad_progress(sns_worker):
with pytest.raises(ProgressTypeError):
create_job(jobid='1', parts=parts, progress=Exception())
sns_worker.assert_not_called()