def test_broker_dequeue_called(self):
item = {
"version": 1,
"task_options": {
"task_id": "f5ceee05-5e41-4fc4-8e2e-d16aa6d67bff",
"eta": "2019-03-24T16:00:12.247687+00:00",
"current_retries": 0,
"max_retries": 0,
"hook_metadata": "",
"args": [],
"kwargs": {
"a": 21,
"b": 535
},
},
}
with mock.patch(
"{broker_path}.dequeue".format(broker_path=self.broker_path()),
new=AsyncMock()) as mock_dequeue:
mock_dequeue.mock.return_value = json.dumps(item)
worker = wiji.Worker(the_task=self.myTask, worker_id="myWorkerID1")
# queue and consume task
self.myTask.synchronous_delay(a=21, b=535)
dequeued_item = self._run(worker.consume_tasks(TESTING=True))
self.assertEqual(dequeued_item["version"], 1)
self.assertTrue(mock_dequeue.mock.called)
self.assertEqual(mock_dequeue.mock.call_args[1],
{"queue_name": self.myTask.queue_name})
def test_broker_done_called(self):
kwargs = {"a": 263_342, "b": 832_429}
with mock.patch(
"{broker_path}.done".format(broker_path=self.broker_path()),
new=AsyncMock()) as mock_broker_done:
worker = wiji.Worker(the_task=self.myTask, worker_id="myWorkerID1")
self.myTask.synchronous_delay(a=kwargs["a"], b=kwargs["b"])
dequeued_item = self._run(worker.consume_tasks(TESTING=True))
self.assertEqual(dequeued_item["version"], 1)
self.assertTrue(mock_broker_done.mock.called)
self.assertEqual(
json.loads(mock_broker_done.mock.call_args[1]["item"])
["task_options"]["kwargs"],
kwargs,
)
self.assertEqual(mock_broker_done.mock.call_args[1]["queue_name"],
self.myTask.queue_name)
self.assertEqual(mock_broker_done.mock.call_args[1]["state"],
wiji.task.TaskState.EXECUTED)
task_options = json.loads(
mock_broker_done.mock.call_args[1]["item"])["task_options"]
self.assertEqual(task_options["kwargs"], kwargs)
self.assertIsNotNone(task_options["task_id"])
self.assertEqual(len(task_options["task_id"]), 36) # len of uuid4
def test_ratelimit_called(self):
with mock.patch("wiji.ratelimiter.SimpleRateLimiter.limit",
new=AsyncMock()) as mock_ratelimit:
worker = wiji.Worker(the_task=self.myTask, worker_id="myWorkerID1")
# queue and consume task
self.myTask.synchronous_delay(a=21, b=535)
dequeued_item = self._run(worker.consume_tasks(TESTING=True))
self.assertEqual(dequeued_item["version"], 1)
self.assertTrue(mock_ratelimit.mock.called)
def setUp(self):
class AdderTask(wiji.task.Task):
the_broker = wiji.broker.InMemoryBroker()
queue_name = "{0}-AdderTaskQueue".format(uuid.uuid4())
drain_duration = 0.01
async def run(self, a, b):
return a + b
_worker = wiji.Worker(the_task=AdderTask(), worker_id="myWorkerID1")
self.workers = [_worker]
self.logger = wiji.logger.SimpleLogger("wiji.TestCliSigHandling")
def test_run_task_called(self):
kwargs = {"a": 263_342, "b": 832_429}
with mock.patch("wiji.worker.Worker.run_task",
new=AsyncMock()) as mock_run_task:
worker = wiji.Worker(the_task=self.myTask, worker_id="myWorkerID1")
self.myTask.synchronous_delay(a=kwargs["a"], b=kwargs["b"])
dequeued_item = self._run(worker.consume_tasks(TESTING=True))
self.assertEqual(dequeued_item["version"], 1)
self.assertTrue(mock_run_task.mock.called)
self.assertEqual(mock_run_task.mock.call_args[1]["a"], kwargs["a"])
self.assertEqual(mock_run_task.mock.call_args[1]["b"], kwargs["b"])
def test_broker_check_called(self):
with mock.patch("wiji.task.Task._broker_check",
new=AsyncMock()) as mock_broker_check:
worker = wiji.Worker(the_task=self.myTask, worker_id="myWorkerID1")
# queue and consume task
self.myTask.synchronous_delay(a=21, b=535)
dequeued_item = self._run(worker.consume_tasks(TESTING=True))
self.assertEqual(dequeued_item["version"], 1)
self.assertTrue(mock_broker_check.mock.called)
self.assertEqual(mock_broker_check.mock.call_args[1],
{"from_worker": True})
async def async_main(logger: wiji.logger.BaseLogger, app_instance: wiji.app.App) -> None:
"""
(i) set signal handlers.
(ii) consume tasks.
(iii) continuously produce watchdog tasks.
"""
watchdog_worker = wiji.Worker(
the_task=wiji.task.WatchDogTask,
use_watchdog=True,
watchdog_duration=app_instance.watchdog_duration,
)
workers = [watchdog_worker]
watch_dog_producer = [
utils._producer.produce_tasks_continously(
task=wiji.task.WatchDogTask, watchdog_duration=app_instance.watchdog_duration
)
]
_queue_names: typing.List[str] = []
for task_class in app_instance.task_classes:
if task_class.queue_name in _queue_names:
# queue names should be unique
raise ValueError(
"There already exists a task with queue_name: {0}".format(task_class.queue_name)
)
_queue_names.append(task_class.queue_name)
task = task_class()
_worker = wiji.Worker(the_task=task)
workers.append(_worker)
del _queue_names
consumers = []
for i in workers:
consumers.append(i.consume_tasks())
gather_tasks = asyncio.gather(
*consumers, *watch_dog_producer, utils.sig._signal_handling(logger=logger, workers=workers)
)
await gather_tasks
def test_task_no_chain(self):
"""
test task with NO chain does not call task.delay
"""
kwargs = {"a": 400, "b": 901}
worker = wiji.Worker(the_task=self.myTask, worker_id="myWorkerID1")
self.myTask.synchronous_delay(a=kwargs["a"], b=kwargs["b"])
with mock.patch("wiji.task.Task.delay",
new=AsyncMock()) as mock_task_delay:
dequeued_item = self._run(worker.consume_tasks(TESTING=True))
self.assertEqual(dequeued_item["version"], 1)
self.assertFalse(mock_task_delay.mock.called)
def test_eta_respected(self):
kwargs = {"a": 7121, "b": 6122}
# NO re-queuing is carried out
worker = wiji.Worker(the_task=self.myTask, worker_id="myWorkerID1")
self.myTask.synchronous_delay(a=kwargs["a"], b=kwargs["b"])
with mock.patch("wiji.task.Task.delay",
new=AsyncMock()) as mock_task_delay:
dequeued_item = self._run(worker.consume_tasks(TESTING=True))
self.assertEqual(dequeued_item["version"], 1)
self.assertFalse(mock_task_delay.mock.called)
# re-queuing is carried out
worker = wiji.Worker(the_task=self.myTask, worker_id="myWorkerID1")
self.myTask.synchronous_delay(
a=kwargs["a"],
b=kwargs["b"],
tas_options=wiji.task.TaskOptions(eta=788.99))
with mock.patch("wiji.task.Task.delay",
new=AsyncMock()) as mock_task_delay:
dequeued_item = self._run(worker.consume_tasks(TESTING=True))
self.assertEqual(dequeued_item["version"], 1)
self.assertTrue(mock_task_delay.mock.called)
self.assertEqual(mock_task_delay.mock.call_args[1], kwargs)
def test_writing_wiji_tests(self):
with mock.patch.object(
AdderTask, "the_broker", wiji.broker.InMemoryBroker()
) as mock_broker:
self.assertDictEqual(DATABASE, {})
view = ExampleView()
view.post(request={"a": 45, "b": 46})
self.assertIsInstance(AdderTask.the_broker, wiji.broker.InMemoryBroker)
self.assertNotIsInstance(AdderTask.the_broker, ExampleRedisBroker)
self.assertEqual(mock_broker._llen(AdderTask.queue_name), 1)
worker = wiji.Worker(the_task=AdderTask())
self._run(worker.consume_tasks(TESTING=True))
self.assertDictEqual(DATABASE, {"result": 91})
def test_no_chaining_if_retrying(self):
"""
test that if parent task is been retried, the chained task is not queued
"""
class DividerTask(wiji.task.Task):
the_broker = self.BROKER
queue_name = "{0}-DividerTaskQueue".format(uuid.uuid4())
async def run(self, a):
res = a / 3
print("divider res: ", res)
return res
class AdderTask(wiji.task.Task):
the_broker = self.BROKER
queue_name = "{0}-AdderTaskQueue".format(uuid.uuid4())
chain = DividerTask
async def run(self, a, b):
res = a + b
await self.retry(
a=221,
b=555,
task_options=wiji.task.TaskOptions(max_retries=2))
return res
MYAdderTask = AdderTask()
kwargs = {"a": 400, "b": 603}
worker = wiji.Worker(the_task=MYAdderTask, worker_id="myWorkerID1")
MYAdderTask.synchronous_delay(a=kwargs["a"], b=kwargs["b"])
with mock.patch.object(
AdderTask, "delay",
new=AsyncMock()) as mock_adder_delay, mock.patch.object(
DividerTask, "delay",
new=AsyncMock()) as mock_divider_delay:
mock_adder_delay.mock.return_value = None
mock_divider_delay.return_value = None
dequeued_item = self._run(worker.consume_tasks(TESTING=True))
self.assertEqual(dequeued_item["version"], 1)
# divider chain is not queued
self.assertFalse(mock_divider_delay.mock.called)
# but adder task is queued again
self.assertTrue(mock_adder_delay.mock.called)
def test_broker_shutdown_called(self):
with mock.patch("{broker_path}.shutdown".format(
broker_path=self.broker_path()),
new=AsyncMock()) as mock_broker_shutdown:
class AdderTask(wiji.task.Task):
the_broker = self.BROKER
# we want this task to be processed slowly
the_ratelimiter = wiji.ratelimiter.SimpleRateLimiter(
execution_rate=1.0)
queue_name = "{0}-TestWorker.test_shutdown".format(
uuid.uuid4())
drain_duration = 1.0
async def run(self, a, b):
res = a + b
return res
_myTask = AdderTask()
worker = wiji.Worker(the_task=_myTask, worker_id="myWorkerID1")
self.assertFalse(worker.SUCCESFULLY_SHUT_DOWN)
# queue a lot of tasks
queued = []
for i in range(1, 20):
_myTask.synchronous_delay(a=9001, b=i)
queued.append(i)
async def call_worker_shutdown():
"""
sleep for a few seconds so that some tasks can be consumed,
then shutdown worker
"""
await asyncio.sleep(5)
await worker.shutdown()
loop = asyncio.get_event_loop()
tasks = asyncio.gather(worker.consume_tasks(TESTING=False),
call_worker_shutdown(),
loop=loop)
loop.run_until_complete(tasks)
self.assertTrue(mock_broker_shutdown.mock.called)
self.assertEqual(
mock_broker_shutdown.mock.call_args[1]["queue_name"],
AdderTask.queue_name)
def test_task_with_chain(self):
"""
test task with chain CALLS task.delay
"""
class DividerTask(wiji.task.Task):
the_broker = self.BROKER
queue_name = "{0}-DividerTaskQueue".format(uuid.uuid4())
async def run(self, a):
res = a / 3
print("divider res: ", res)
return res
class AdderTask(wiji.task.Task):
the_broker = self.BROKER
queue_name = "{0}-AdderTaskQueue".format(uuid.uuid4())
chain = DividerTask
async def run(self, a, b):
res = a + b
return res
MYAdderTask = AdderTask()
kwargs = {"a": 400, "b": 603}
worker = wiji.Worker(the_task=MYAdderTask, worker_id="myWorkerID1")
MYAdderTask.synchronous_delay(a=kwargs["a"], b=kwargs["b"])
with mock.patch.object(
AdderTask, "delay",
new=AsyncMock()) as mock_adder_delay, mock.patch.object(
DividerTask, "delay",
new=AsyncMock()) as mock_divider_delay:
mock_adder_delay.mock.return_value = None
mock_divider_delay.return_value = None
dequeued_item = self._run(worker.consume_tasks(TESTING=True))
self.assertEqual(dequeued_item["version"], 1)
# adder task is not queued
self.assertFalse(mock_adder_delay.mock.called)
# but divider task(the chain) is queued
self.assertTrue(mock_divider_delay.mock.called)
self.assertEqual(mock_divider_delay.mock.call_args[0][1],
kwargs["a"] + kwargs["b"])
def test_no_chaining_if_exception(self):
"""
test that if parent task raises exception, the chained task is not queued
"""
class DividerTask(wiji.task.Task):
the_broker = self.BROKER
queue_name = "{0}-DividerTaskQueue".format(uuid.uuid4())
async def run(self, a):
res = a / 3
print("divider res: ", res)
return res
class AdderTask(wiji.task.Task):
the_broker = self.BROKER
queue_name = "{0}-AdderTaskQueue".format(uuid.uuid4())
async def run(self, a, b):
return await self.do_work(a, b)
@staticmethod
async def do_work(a, b):
return a + b
MYAdderTask = AdderTask()
kwargs = {"a": 400, "b": 603}
worker = wiji.Worker(the_task=MYAdderTask, worker_id="myWorkerID1")
MYAdderTask.synchronous_delay(a=kwargs["a"], b=kwargs["b"])
with mock.patch(
"wiji.task.Task.delay",
new=AsyncMock()) as mock_task_delay, mock.patch.object(
AdderTask,
"do_work",
side_effect=Exception(
"test_no_chaining_if_exception")) as mock_do_work:
mock_task_delay.mock.return_value = None
_ = mock_do_work
dequeued_item = self._run(worker.consume_tasks(TESTING=True))
self.assertEqual(dequeued_item["version"], 1)
# chain is not queued
self.assertFalse(mock_task_delay.mock.called)
def test_shutdown(self):
class AdderTask(wiji.task.Task):
the_broker = self.BROKER
# we want this task to be processed slowly
the_ratelimiter = wiji.ratelimiter.SimpleRateLimiter(
execution_rate=1.0)
queue_name = "{0}-TestWorker.test_shutdown".format(uuid.uuid4())
drain_duration = 1.0
async def run(self, a, b):
res = a + b
return res
_myTask = AdderTask()
worker = wiji.Worker(the_task=_myTask, worker_id="myWorkerID1")
self.assertFalse(worker.SUCCESFULLY_SHUT_DOWN)
# queue a lot of tasks
queued = []
for i in range(1, 20):
_myTask.synchronous_delay(a=9001, b=i)
queued.append(i)
async def call_worker_shutdown():
"""
sleep for a few seconds so that some tasks can be consumed,
then shutdown worker
"""
await asyncio.sleep(5)
await worker.shutdown()
loop = asyncio.get_event_loop()
tasks = asyncio.gather(worker.consume_tasks(TESTING=False),
call_worker_shutdown(),
loop=loop)
loop.run_until_complete(tasks)
# assert that some tasks have been consumed and also
# that not all were consumed.
self.assertTrue(_myTask.the_broker._llen(AdderTask.queue_name) > 10)
self.assertTrue(
_myTask.the_broker._llen(AdderTask.queue_name) < len(queued))
self.assertTrue(worker.SUCCESFULLY_SHUT_DOWN)
def test_memory_leak(self):
"""
This guards against a regression.
wiji.broker.InMemoryBroker had a memory leak: https://github.com/komuw/wiji/issues/71
"""
class AdderTask(wiji.task.Task):
the_broker = self.BROKER
queue_name = "{0}-TestWorker.test_cool".format(uuid.uuid4())
async def run(self, a, b):
res = a + b
return res
_myTask = AdderTask()
worker = wiji.Worker(the_task=_myTask, worker_id="myWorkerID1")
_myTask.synchronous_delay(a=9001, b=6)
self.assertEqual(_myTask.the_broker._llen(AdderTask.queue_name), 1)
self._run(worker.consume_tasks(TESTING=True))
self.assertEqual(_myTask.the_broker._llen(AdderTask.queue_name), 0)
def test_no_rlimit(self):
with mock.patch("wiji.hook.SimpleHook.notify",
new=AsyncMock()) as mock_hook_notify:
mock_hook_notify.mock.return_value = None
worker = wiji.Worker(the_task=self.myAdderTask,
worker_id="myWorkerID1")
# queue task
kwargs = {"a": 78, "b": 101}
self.myAdderTask.synchronous_delay(a=kwargs["a"], b=kwargs["b"])
# consume
dequeued_item = self._run(worker.consume_tasks(TESTING=True))
self.assertEqual(dequeued_item["task_options"]["kwargs"], kwargs)
self.assertTrue(mock_hook_notify.mock.called)
self.assertEqual(
mock_hook_notify.mock.call_args[1]["return_value"],
kwargs["a"] + kwargs["b"])
self.assertEqual(
mock_hook_notify.mock.call_args[1]["execution_exception"],
None)
def test_consume_tasks(self):
worker = wiji.Worker(the_task=self.myTask, worker_id="myWorkerID1")
# queue task
kwargs = {"a": 78, "b": 101}
self.myTask.synchronous_delay(a=kwargs["a"], b=kwargs["b"])
# consume
dequeued_item = self._run(worker.consume_tasks(TESTING=True))
self.assertEqual(dequeued_item["version"], 1)
self.assertEqual(dequeued_item["task_options"]["current_retries"], 0)
self.assertEqual(dequeued_item["task_options"]["max_retries"], 0)
self.assertEqual(dequeued_item["task_options"]["args"], [])
self.assertEqual(dequeued_item["task_options"]["kwargs"], kwargs)
# queue task
self.myTask.synchronous_delay(34, 88)
# consume
dequeued_item = self._run(worker.consume_tasks(TESTING=True))
self.assertEqual(dequeued_item["version"], 1)
self.assertEqual(dequeued_item["task_options"]["current_retries"], 0)
self.assertEqual(dequeued_item["task_options"]["max_retries"], 0)
self.assertEqual(dequeued_item["task_options"]["args"], [34, 88])
self.assertEqual(dequeued_item["task_options"]["kwargs"], {})
def test_success_instantiation(self):
wiji.Worker(the_task=self.myTask, worker_id="myWorkerID1")
def mock_create_worker():
wiji.Worker(the_task=self.myTask, worker_id=92033)
def mock_create_worker():
wiji.Worker(the_task="bad-task-arg")