Exemplo n.º 1
0
    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})
Exemplo n.º 2
0
    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
Exemplo n.º 3
0
    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)
Exemplo n.º 4
0
    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")
Exemplo n.º 5
0
    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"])
Exemplo n.º 6
0
    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})
Exemplo n.º 7
0
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
Exemplo n.º 8
0
    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)
Exemplo n.º 9
0
    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)
Exemplo n.º 10
0
    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})
Exemplo n.º 11
0
    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)
Exemplo n.º 12
0
    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)
Exemplo n.º 13
0
    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"])
Exemplo n.º 14
0
    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)
Exemplo n.º 15
0
    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)
Exemplo n.º 16
0
    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)
Exemplo n.º 17
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)
Exemplo n.º 18
0
    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"], {})
Exemplo n.º 19
0
 def test_success_instantiation(self):
     wiji.Worker(the_task=self.myTask, worker_id="myWorkerID1")
Exemplo n.º 20
0
 def mock_create_worker():
     wiji.Worker(the_task=self.myTask, worker_id=92033)
Exemplo n.º 21
0
 def mock_create_worker():
     wiji.Worker(the_task="bad-task-arg")