def test_status(parts, monkeypatch):
"""New job shows all parts remaining
"""
monkeypatch.setenv('WorkTopic', 'abc123')
with patch('redis.StrictRedis', mock_strict_redis_client):
p = RedisProgress()
p.set_total('123', parts)
status = p.status('123')
assert status['total'] == 3
assert status['remaining'] == 3
assert status['progress'] == 0
def test_no_delete_when_done(parts, monkeypatch):
p = RedisProgress(host='localhost', port=6379, db=0, topic_arn='nope', delete_when_done=False)
jobid = '123'
p.set_total(jobid, [parts[0]])
assert len(list(p.list_pending_parts(jobid))) == 1
p.complete_part(jobid, 0)
assert 'total' in p._decode_dict(p.redis.hgetall('123-metadata'))
assert p.status(jobid)['remaining'] == 0 # status still works
def test_delete(parts, monkeypatch):
p = RedisProgress(host='localhost', port=6379, db=0, topic_arn='nope')
jobid = '123'
p.set_total(jobid, parts)
assert 'total' in p.status(jobid)
p.delete(jobid)
with pytest.raises(JobDoesNotExist):
p.status(jobid)
def test_list_pending_nojobs(parts, monkeypatch):
p = RedisProgress(host='localhost', port=6379, db=0, topic_arn='nope')
jobid = '123'
p.set_total(jobid, parts)
assert len(list(p.list_pending_parts(jobid))) == 3
p.delete(jobid)
with pytest.raises(JobDoesNotExist):
p.list_pending_parts(jobid)
def test_delete_when_done(parts, monkeypatch):
p = RedisProgress(host='localhost', port=6379, db=0, topic_arn='nope', delete_when_done=True)
jobid = '123'
p.set_total(jobid, [parts[0]])
assert len(list(p.list_pending_parts(jobid))) == 1
p.complete_part(jobid, 0)
assert not p.redis.hgetall('123-metadata')
with pytest.raises(JobDoesNotExist):
p.status(jobid)
def test_list_pending(parts, monkeypatch):
p = RedisProgress(host='localhost', port=6379, db=0, topic_arn='nope')
jobid = '123'
p.set_total(jobid, parts)
assert len(list(p.list_pending_parts(jobid))) == 3
p.complete_part(jobid, 0)
assert len(list(p.list_pending_parts(jobid))) == 2
def test_status_no_total(monkeypatch):
""" Have not created a job with .set_total() yet
"""
monkeypatch.setenv('WorkTopic', 'abc123')
with patch('redis.StrictRedis', mock_strict_redis_client):
with pytest.raises(JobDoesNotExist):
RedisProgress().status('123')
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_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_metadata(parts, monkeypatch):
"""Setting and getting job metadata works
"""
monkeypatch.setenv('WorkTopic', 'abc123')
with patch('redis.StrictRedis', mock_strict_redis_client):
p = RedisProgress()
p.set_total('123', parts)
p.set_metadata('123', {'test': 'foo'})
assert p.status('123')['metadata']['test'] == 'foo'
def test_failjob(parts, monkeypatch):
"""Mark job as failed works
"""
monkeypatch.setenv('WorkTopic', 'abc123')
with patch('redis.StrictRedis', mock_strict_redis_client):
p = RedisProgress()
p.set_total('123', parts)
p.fail_job('123', 'epic fail')
assert p.status('123')['failed'] is True
def test_status_part(parts, monkeypatch):
"""Check if part is complete
"""
monkeypatch.setenv('WorkTopic', 'abc123')
with patch('redis.StrictRedis', mock_strict_redis_client):
p = RedisProgress()
p.set_total('123', parts)
p.complete_part('123', 0)
assert p.status('123', part=0)['complete'] is True
assert p.status('123', part=1)['complete'] is False
def test_status_complete_some(parts, monkeypatch):
"""job shows partial progress
"""
monkeypatch.setenv('WorkTopic', 'abc123')
with patch('redis.StrictRedis', mock_strict_redis_client):
p = RedisProgress()
p.set_total('123', parts)
done_yet = p.complete_part('123', 0)
assert not done_yet
status = p.status('123')
assert status['total'] == 3
assert status['remaining'] == 2
assert status['progress'] == 1 / 3
def test_status_complete_all(parts, monkeypatch):
"""job shows complete progress
"""
monkeypatch.setenv('WorkTopic', 'abc123')
with patch('redis.StrictRedis', mock_strict_redis_client):
p = RedisProgress()
p.set_total('123', parts)
for i, _ in enumerate(parts):
done_yet = p.complete_part('123', i)
assert done_yet
status = p.status('123')
assert status['total'] == 3
assert status['remaining'] == 0
assert status['progress'] == 1.0
def validate_db(ctx, param, value):
topic_arn = 'should-never-need-to-use-SNS'
if value.startswith('redis'):
url = urlparse(value)
assert url.scheme == 'redis'
db = 0
if url.query:
queries = url.query.split('&')
for query in queries:
k, v = query.split('=')
if k == 'db':
db = int(v)
host, port = url.netloc.split(":")
return RedisProgress(host=host, port=port, db=db, topic_arn=topic_arn)
elif value.startswith('arn:'):
return DynamoProgress(table_arn=value, topic_arn=topic_arn)
else:
raise click.BadParameter('Needs to be {}'.format(DBHELP))
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_list_jobids(parts, monkeypatch):
p = RedisProgress(host='localhost', port=6379, db=0, topic_arn='nope')
p.set_total('job1', parts)
p.set_total('job2', parts)
assert list(p.list_jobs(status=False)) == ['job1', 'job2']
from . import version as app_version
from cogeo_mosaic.utils import (
_aws_get_data,
_aws_put_data,
_compress_gz_json,
create_mosaic,
get_hash,
)
logger = logging.getLogger("watchbot")
logger.setLevel(logging.INFO)
progress = RedisProgress(
host=os.environ["REDIS_DB_HOST"],
port=os.environ["REDIS_DB_PORT"],
topic_arn=os.environ["SNS_TOPIC"],
)
class LambdaFileIsTooBig(Exception):
"""File is too big for AWS Lambda."""
def _aws_list_objects(bucket, prefix):
"""List object in a bucket."""
session = boto3_session()
client = session.client("s3")
pag = client.get_paginator("list_objects_v2")
files = []
def test_list_jobs_meta(parts, monkeypatch):
p = RedisProgress(host='localhost', port=6379, db=0, topic_arn='nope')
p.set_total('job1', parts)
p.set_total('job2', parts)
p.redis.delete('job1-parts') # job1-metdata should remain thus still be in list
assert len(list(p.list_jobs())) == 2
def test_list_jobs(parts, monkeypatch):
p = RedisProgress(host='localhost', port=6379, db=0, topic_arn='nope')
p.set_total('job1', parts)
p.set_total('job2', parts)
assert len(list(p.list_jobs())) == 2