def test_all():
w1 = Worker('w1', queues=[QueueFactory()])
w2 = Worker('w2', queues=[QueueFactory()])
assert Worker.all() == []
w1.startup()
w2.startup()
assert id_list(Worker.all()) == [w1.id, w2.id]
w1.shutdown()
assert id_list(Worker.all()) == [w2.id]
def test_died(time_mocker, connection, assert_atomic):
time = time_mocker('redis_tasks.worker.utcnow')
# Die while idle
worker = Worker('idleworker', queues=[QueueFactory()])
time.step()
worker.startup()
time.step()
assert connection.ttl(worker.key) == -1
assert worker.id in worker_registry.get_worker_ids()
with assert_atomic():
worker.died()
assert worker.id not in worker_registry.get_worker_ids()
for w in [worker, Worker.fetch(worker.id)]:
assert w.state == WorkerState.DEAD
assert w.shutdown_at == time.now
assert connection.ttl(worker.key) > 0
# die whith task in limbo
worker = Worker('limboworker', queues=[QueueFactory(), QueueFactory()])
queue = worker.queues[1]
time.step()
worker.startup()
time.step()
queue.enqueue_call()
task = queue.dequeue(worker)
with assert_atomic(exceptions=['hgetall']):
worker.died()
assert queue.get_task_ids() == [task.id]
assert worker.id not in worker_registry.get_worker_ids()
for w in [worker, Worker.fetch(worker.id)]:
assert w.state == WorkerState.DEAD
assert w.current_task_id is None
assert w.shutdown_at == time.now
assert connection.ttl(worker.key) > 0
# die while busy
worker = Worker('busyworker', queues=[QueueFactory(), QueueFactory()])
queue = worker.queues[1]
time.step()
worker.startup()
time.step()
queue.enqueue_call()
task = queue.dequeue(worker)
worker.start_task(task)
with assert_atomic(exceptions=['hgetall']):
worker.died()
assert queue.get_task_ids() == []
assert failed_task_registry.get_task_ids() == [task.id]
assert worker.id not in worker_registry.get_worker_ids()
for w in [worker, Worker.fetch(worker.id)]:
assert w.state == WorkerState.DEAD
assert w.current_task_id is None
assert w.shutdown_at == time.now
assert connection.ttl(worker.key) > 0
def test_queue_registry(assert_atomic, connection):
registry = registries.queue_registry
queue1 = QueueFactory()
queue2 = QueueFactory()
with assert_atomic():
registry.add(queue1)
registry.add(queue2)
assert set(registry.get_names()) == {queue1.name, queue2.name}
with assert_atomic():
registry.remove(queue1)
assert registry.get_names() == [queue2.name]
registry.remove(queue2)
assert registry.get_names() == []
def test_fetch_current_task():
worker = Worker('testworker', queues=[QueueFactory()])
queue = worker.queues[0]
queue.enqueue_call()
assert worker.fetch_current_task() is None
task = queue.dequeue(worker)
assert worker.fetch_current_task().id == task.id
def test_run_shutdown(settings, mocker):
mocker.patch.object(WorkerProcess, 'queue_iter', side_effect=ShutdownRequested)
wp = WorkerProcess([QueueFactory()])
with pytest.raises(ShutdownRequested):
wp.run(False)
assert wp.worker.started_at
assert wp.worker.shutdown_at
def test_queue_iter(connection, mocker):
queues = [QueueFactory(), QueueFactory()]
wp = WorkerProcess(queues)
wp.worker.startup()
qi = wp.queue_iter(False)
task1 = queues[1].enqueue_call()
task2 = queues[0].enqueue_call()
assert next(qi).id == task2.id
assert next(qi).id == task1.id
def my_await_multi(*args):
nonlocal await_counter, task3, task4, task5
await_counter += 1
if await_counter == 1:
return None
elif await_counter == 2:
task3 = queues[1].enqueue_call()
return queues[1]
elif await_counter == 3:
task4 = queues[1].enqueue_call()
task5 = queues[1].enqueue_call()
return queues[0]
assert False
await_counter = 0
task3 = task4 = task5 = None
mocker.patch('redis_tasks.queue.Queue.await_multi', new=my_await_multi)
assert next(qi).id == task3.id
assert wp.worker.current_task_id == task3.id
assert await_counter == 2
assert next(qi).id == task4.id
assert wp.worker.current_task_id == task4.id
assert await_counter == 3
assert next(qi).id == task5.id
assert await_counter == 3
# Assert that tasks are moved to the worker
assert (decode_list(connection.lrange(wp.worker.task_key, 0, -1)) ==
id_list([task5, task4, task3, task1, task2]))
# Burst mode
wp = WorkerProcess(queues)
wp.worker.startup()
task1 = queues[1].enqueue_call()
task2 = queues[0].enqueue_call()
task3 = queues[1].enqueue_call()
assert id_list(wp.queue_iter(True)) == id_list([task2, task1, task3])
def test_signal_shutdown_in_queuewait():
wp = WorkerProcess([QueueFactory()])
process = multiprocessing.Process(target=wp.run)
process.start()
time.sleep(0.1)
os.kill(process.pid, signal.SIGTERM)
process.join(1)
assert not process.is_alive()
def test_state_transitions(time_mocker, connection, assert_atomic):
worker = Worker('myworker', queues=[QueueFactory(), QueueFactory()])
time = time_mocker('redis_tasks.worker.utcnow')
queue = worker.queues[0]
time.step()
assert not connection.exists(worker.key, worker.task_key)
assert worker.id not in worker_registry.get_worker_ids()
with assert_atomic():
worker.startup()
assert worker.id in worker_registry.get_worker_ids()
for w in [worker, Worker.fetch(worker.id)]:
assert w.started_at == time.now
assert w.state == WorkerState.IDLE
queue.enqueue_call()
with assert_atomic(exceptions=['hgetall']):
task = queue.dequeue(worker)
assert connection.exists(worker.key, worker.task_key) == 2
for w in [worker, Worker.fetch(worker.id)]:
assert w.current_task_id == task.id
assert w.state == WorkerState.IDLE
with assert_atomic():
worker.start_task(task)
for w in [worker, Worker.fetch(worker.id)]:
assert w.current_task_id == task.id
assert w.state == WorkerState.BUSY
with assert_atomic():
worker.end_task(task, TaskOutcome("success"))
for w in [worker, Worker.fetch(worker.id)]:
assert w.current_task_id is None
assert w.state == WorkerState.IDLE
time.step()
assert connection.ttl(worker.key) == -1
assert worker.id in worker_registry.get_worker_ids()
with assert_atomic():
worker.shutdown()
assert worker.id not in worker_registry.get_worker_ids()
for w in [worker, Worker.fetch(worker.id)]:
assert w.state == WorkerState.DEAD
assert w.shutdown_at == time.now
assert connection.ttl(worker.key) > 0
def test_signal_shutdown_in_task(suprocess_socket):
queue = QueueFactory()
task = queue.enqueue_call(taskwait)
wp = WorkerProcess([queue])
process = multiprocessing.Process(target=wp.run)
process.start()
with suprocess_socket.accept() as taskconn:
assert taskconn.poll(1)
assert taskconn.recv() == "A"
os.kill(process.pid, signal.SIGTERM)
assert taskconn.poll(1)
assert taskconn.recv() == "B"
process.join(1)
assert not process.is_alive()
task.refresh()
assert task.status == TaskStatus.FAILED
assert 'Worker shutdown' in task.error_message.splitlines()[-1]
def test_info(cli_run, stub):
cli_run('info')
for i in range(2):
QueueFactory().enqueue_call(stub)
for i in range(3):
WorkerFactory().startup()
result = cli_run('info')
assert "2 queue(s), 2 task(s) total" in result.output
assert "3 worker(s)" in result.output
def test_heartbeats(mocker, stub):
heartbeat_sent = False
def send_heartbeat(*args, **kwargs):
nonlocal heartbeat_sent
heartbeat_sent = True
mocker.patch('redis_tasks.registries.WorkerRegistry.heartbeat', side_effect=send_heartbeat)
def consume_heartbeat(*args, **kwargs):
nonlocal heartbeat_sent
assert heartbeat_sent is True
heartbeat_sent = False
return mocker.DEFAULT
maintenance = mocker.patch('redis_tasks.worker_process.Maintenance.run_if_neccessary',
side_effect=consume_heartbeat)
def my_await(*args):
nonlocal task
for i in range(3):
consume_heartbeat()
print("Awaited!")
yield None
task = queue.enqueue_call(stub)
consume_heartbeat()
yield queue
raise ShutdownRequested()
mock_await = mocker.patch('redis_tasks.queue.Queue.await_multi', side_effect=my_await())
mocker.patch.object(WorkHorse, 'start', new=WorkHorse.run)
horse_alive = mocker.patch.object(WorkHorse, 'is_alive', return_value=True)
def my_join(*args):
consume_heartbeat()
yield None
consume_heartbeat()
yield None
consume_heartbeat()
horse_alive.return_value = False
yield None
mock_join = mocker.patch.object(WorkHorse, 'join', side_effect=my_join())
stub.mock.reset_mock()
queue = QueueFactory()
task = None
wp = WorkerProcess([queue])
with pytest.raises(ShutdownRequested):
wp.run(False)
task.refresh()
assert task.status == TaskStatus.FINISHED
assert maintenance.call_count == 2
assert mock_await.call_count == 5
assert mock_join.call_count == 3
assert stub.mock.called
def test_run(settings, mocker, stub):
tasks = [mocker.sentinel.t1, mocker.sentinel.t2, mocker.sentinel.t3]
mocker.patch('redis_tasks.worker_process.WorkerProcess.queue_iter', return_value=tasks)
def my_process(task):
assert wp.worker.state == WorkerState.IDLE
process = mocker.patch('redis_tasks.worker_process.WorkerProcess.process_task',
side_effect=my_process)
maintenance = mocker.patch('redis_tasks.worker_process.Maintenance')
wp = WorkerProcess([QueueFactory()])
assert wp.run(True) == 3
assert process.call_args_list == [((x, ),) for x in tasks]
assert maintenance().run_if_neccessary.call_count == 4
assert wp.worker.state == WorkerState.DEAD
settings.WORKER_PRELOAD_FUNCTION = stub.path
stub.mock.reset_mock()
wp = WorkerProcess([QueueFactory()])
wp.run(True)
assert stub.mock.call_count == 1
def test_cancel(assert_atomic, connection):
q = QueueFactory()
task = q.enqueue_call()
with assert_atomic():
task.cancel()
assert q.get_task_ids() == []
assert task.status == TaskStatus.CANCELED
assert not connection.exists(task.key)
w = WorkerFactory()
w.startup()
task = q.enqueue_call()
q.dequeue(w)
with pytest.raises(InvalidOperation):
with assert_atomic():
task.cancel()
assert worker_registry.get_running_tasks() == {w.id: task.id}
assert connection.exists(task.key)
def test_process_task(mocker):
q = QueueFactory()
wp = WorkerProcess([q])
wp.worker.startup()
q.enqueue_call()
task = q.dequeue(wp.worker)
def my_execute(task):
assert wp.worker.state == WorkerState.BUSY
assert task.status == TaskStatus.RUNNING
return TaskOutcome("success")
execute = mocker.patch.object(WorkerProcess, 'execute_task', side_effect=my_execute)
wp.process_task(task)
assert task.status == TaskStatus.FINISHED
q.enqueue_call()
task = q.dequeue(wp.worker)
execute.side_effect = ArithmeticError()
wp.process_task(task)
assert task.status == TaskStatus.FAILED
assert 'ArithmeticError' in task.error_message
def test_empty(cli_run, stub):
queues = [QueueFactory() for i in range(5)]
for q in queues:
q.enqueue_call(stub)
with pytest.raises(click.UsageError):
cli_run('empty')
assert all(q.count() == 1 for q in queues)
cli_run('empty', queues[0].name, queues[1].name)
assert queues[0].count() == 0
assert queues[1].count() == 0
assert queues[2].count() == 1
assert set(queues) == set(Queue.all())
cli_run('empty', '--delete', queues[1].name, queues[2].name)
assert set(queues) - {queues[1], queues[2]} == set(Queue.all())
cli_run('empty', '--all')
assert all(q.count() == 0 for q in queues)
def test_worker_death(assert_atomic, connection, stub):
def setup(func):
w = WorkerFactory()
w.startup()
q.enqueue_call(func)
task = q.dequeue(w)
return task, w
q = QueueFactory()
# Worker died before starting work on the task
task, w = setup(stub)
assert worker_registry.get_running_tasks() == {w.id: task.id}
with assert_atomic(exceptions=['hgetall']):
w.died()
assert q.get_task_ids() == [task.id]
assert worker_registry.get_running_tasks() == dict()
assert failed_task_registry.get_task_ids() == []
# Worker died after starting non-reentrant task
q.empty()
task, w = setup(stub)
assert worker_registry.get_running_tasks() == {w.id: task.id}
w.start_task(task)
with assert_atomic(exceptions=['hgetall']):
w.died()
assert q.get_task_ids() == []
assert worker_registry.get_running_tasks() == dict()
assert failed_task_registry.get_task_ids() == [task.id]
# Worker died after starting reentrant task
connection.delete(failed_task_registry.key)
task, w = setup(reentrant_stub)
assert worker_registry.get_running_tasks() == {w.id: task.id}
with assert_atomic():
w.start_task(task)
with assert_atomic(exceptions=['hgetall']):
w.died()
assert q.get_task_ids() == [task.id]
assert worker_registry.get_running_tasks() == dict()
assert failed_task_registry.get_task_ids() == []
def test_state_transistions(assert_atomic, connection, time_mocker, stub):
time = time_mocker('redis_tasks.task.utcnow')
task = Task(reentrant_stub)
q = QueueFactory()
w = WorkerFactory()
w.startup()
# enqueue
time.step()
assert not connection.exists(task.key)
with assert_atomic():
task.enqueue(q)
assert q.get_task_ids() == [task.id]
assert connection.exists(task.key)
for t in [task, Task.fetch(task.id)]:
assert t.enqueued_at == time.now
assert t.status == TaskStatus.QUEUED
assert t.origin == q.name
# dequeue
task = q.dequeue(w)
assert q.get_task_ids() == []
assert worker_registry.get_running_tasks() == {w.id: task.id}
# set_running
time.step()
with assert_atomic():
w.start_task(task)
assert worker_registry.get_running_tasks() == {w.id: task.id}
for t in [task, Task.fetch(task.id)]:
assert t.status == TaskStatus.RUNNING
assert t.started_at == time.now
# requeue
time.step()
with assert_atomic():
w.end_task(task, TaskOutcome("requeue"))
assert worker_registry.get_running_tasks() == dict()
assert q.get_task_ids() == [task.id]
for t in [task, Task.fetch(task.id)]:
assert t.status == TaskStatus.QUEUED
assert t.started_at is None
# set_finished
task = q.dequeue(w)
w.start_task(task)
time.step()
with assert_atomic():
w.end_task(task, TaskOutcome("success"))
assert q.get_task_ids() == []
assert finished_task_registry.get_task_ids() == [task.id]
for t in [task, Task.fetch(task.id)]:
assert t.status == TaskStatus.FINISHED
assert t.ended_at == time.now
# set_failed
task = q.enqueue_call(stub)
task = q.dequeue(w)
w.start_task(task)
assert worker_registry.get_running_tasks() == {w.id: task.id}
time.step()
with assert_atomic():
w.end_task(task, TaskOutcome("failure", message="my error"))
assert worker_registry.get_running_tasks() == dict()
assert failed_task_registry.get_task_ids() == [task.id]
for t in [task, Task.fetch(task.id)]:
assert t.status == TaskStatus.FAILED
assert t.error_message == "my error"
assert t.ended_at == time.now
def test_worker_reg_running_tasks():
registry = registries.worker_registry
queue = QueueFactory()
t1 = queue.enqueue_call()
t2 = queue.enqueue_call()
worker1 = WorkerFactory(queues=[queue])
worker2 = WorkerFactory(queues=[queue])
worker1.startup()
worker2.startup()
assert registry.get_running_tasks() == dict()
queue.push(t1)
queue.dequeue(worker1)
assert registry.get_running_tasks() == {worker1.id: t1.id}
worker1.start_task(t1)
queue.push(t2)
queue.dequeue(worker2)
worker2.start_task(t2)
assert registry.get_running_tasks() == {
worker1.id: t1.id,
worker2.id: t2.id
}
worker1.end_task(t1, TaskOutcome("success"))
assert registry.get_running_tasks() == {worker2.id: t2.id}
def test_heartbeat(mocker):
heartbeat = mocker.patch.object(worker_registry, 'heartbeat')
worker = Worker('testworker', queues=[QueueFactory()])
worker.startup()
worker.heartbeat()
assert heartbeat.called_once_with(worker)
def test_execute_task(mocker, settings, time_mocker):
horse = None
def my_start(self):
nonlocal horse
horse = self
assert horse.task == task
horse_alive.return_value = True
horse.worker_connection.send(True)
mocker.patch.object(WorkHorse, 'start', new=my_start)
horse_alive = mocker.patch.object(WorkHorse, 'is_alive', return_value=False)
# Normal run
def my_join():
yield None
yield None
horse.worker_connection.send(TaskOutcome("success"))
horse_alive.return_value = False
yield None
mock_join = mocker.patch.object(WorkHorse, 'join', side_effect=my_join())
wp = WorkerProcess([QueueFactory()])
wp.worker.startup()
task = TaskFactory()
outcome = wp.execute_task(task)
assert outcome.outcome == "success"
assert mock_join.call_count == 3
# Unexpected WorkHorse death
def dying_join():
yield None
horse_alive.return_value = False
yield None
mock_join.side_effect = dying_join()
outcome = wp.execute_task(task)
assert outcome.outcome == "failure"
assert "Workhorse died unexpectedly" in outcome.message
# Shutdown
shutdown_initiated = False
def take_usr1(signum):
nonlocal shutdown_initiated
shutdown_initiated = True
assert signum == signal.SIGUSR1
def shutdown_join():
yield None
assert wp.in_interruptible
yield ShutdownRequested()
assert shutdown_initiated
yield None
horse.worker_connection.send(TaskOutcome("requeue"))
yield None
horse_alive.return_value = False
yield None
mock_join.side_effect = shutdown_join()
fake_signal = mocker.patch.object(WorkHorse, 'send_signal', side_effect=take_usr1)
outcome = wp.execute_task(task)
assert outcome.outcome == "requeue"
# Timeout
def timeout_join():
yield None
yield None
time.step()
yield None
yield None
def take_kill(signum):
assert signum == signal.SIGKILL
horse_alive.return_value = False
mock_join.side_effect = timeout_join()
settings.DEFAULT_TASK_TIMEOUT = 1
time = time_mocker("redis_tasks.worker_process.utcnow")
time.step()
fake_signal = mocker.patch.object(WorkHorse, 'send_signal', side_effect=take_kill)
outcome = wp.execute_task(task)
assert outcome.outcome == "failure"
assert "Task timeout (1 sec) reached" in outcome.message
fake_signal.assert_called_once_with(9)
def test_persistence(assert_atomic, connection, time_mocker):
time = time_mocker('redis_tasks.worker.utcnow')
time.step()
fields = {
'description', 'state', 'queues', 'started_at', 'shutdown_at',
'current_task_id'
}
def randomize_data(worker):
string_fields = ['description', 'state', 'current_task_id']
date_fields = ['started_at', 'shutdown_at']
for f in string_fields:
setattr(worker, f, str(uuid.uuid4()))
for f in date_fields:
setattr(
worker, f,
datetime.datetime(random.randint(1000, 9999),
1,
1,
tzinfo=datetime.timezone.utc))
worker.args = tuple(str(uuid.uuid4()) for i in range(4))
worker.kwargs = {str(uuid.uuid4()): ["d"]}
worker.meta = {"x": [str(uuid.uuid4())]}
worker.aborted_runs = ["foo", "bar", str(uuid.uuid4())]
def as_dict(worker):
return {
f: getattr(worker, f)
if f != 'queues' else [q.name for q in worker.queues]
for f in fields
}
worker = Worker("testworker", queues=[QueueFactory()])
with assert_atomic():
worker.startup()
assert as_dict(worker) == as_dict(Worker.fetch(worker.id))
worker2 = Worker("worker2", queues=[QueueFactory() for i in range(5)])
worker2.startup()
assert as_dict(worker2) == as_dict(Worker.fetch(worker2.id))
assert as_dict(worker) != as_dict(worker2)
worker = Worker("testworker", queues=[QueueFactory()])
worker.startup()
with assert_atomic():
worker._save()
assert set(decode_list(connection.hkeys(worker.key))) <= fields
assert as_dict(Worker.fetch(worker.id)) == as_dict(worker)
randomize_data(worker)
worker._save()
assert as_dict(Worker.fetch(worker.id)) == as_dict(worker)
# only deletes
worker.started_at = None
worker._save(['started_at'])
assert as_dict(Worker.fetch(worker.id)) == as_dict(worker)
for i in range(5):
store = random.sample(fields, 3)
copy = Worker.fetch(worker.id)
randomize_data(copy)
copy._save(store)
for f in store:
setattr(worker, f, getattr(copy, f))
assert as_dict(Worker.fetch(worker.id)) == as_dict(worker)
worker = Worker("nonexist", queues=[QueueFactory()])
with pytest.raises(WorkerDoesNotExist):
worker.refresh()
with pytest.raises(WorkerDoesNotExist):
Worker.fetch("nonexist")