def test_run(self):
worker_task_callback = MagicMock()
worker = _WorkerProcess('testapp', ['testqueue'], 10,
worker_task_callback, None, False)
worker._consume = MagicMock(side_effect=SystemExit)
worker._run()
self.assertFalse(worker_task_callback.called)
worker._consume.assert_called_with()
def test_run_exception(self, print_exc_mock):
worker_task_callback = MagicMock()
worker = _WorkerProcess('testapp', ['testqueue'], 10,
worker_task_callback, None, False)
worker._consume = MagicMock(side_effect=AssertionError)
self.assertRaises(AssertionError, worker._run)
self.assertFalse(worker_task_callback.called)
print_exc_mock.assert_called_with()
worker._consume.assert_called_with()
def test_run_accessrefused(self, sleep_mock):
worker_task_callback = MagicMock()
worker = _WorkerProcess('testapp', ['testqueue'], 10,
worker_task_callback, None, False)
worker._consume = MagicMock(side_effect=[AccessRefused, None])
worker._run()
self.assertFalse(worker_task_callback.called)
sleep_mock.assert_called_once_with(ANY)
worker._consume.assert_called_with()
def test_start_worker_child(self, exit_mock, fork_mock):
app = MagicMock()
worker = _WorkerProcess(app, ['testqueue'], None, None, None, False)
worker._run = MagicMock()
master = WorkerMaster('testapp')
fork_mock.return_value = 0
master._start_worker(worker)
fork_mock.assert_called_with()
worker._run.assert_called_with()
exit_mock.assert_called_with(os.EX_OK)
def test_on_message_task_callback(self):
app = MagicMock()
task_cb = MagicMock(side_effect=SystemExit)
worker = _WorkerProcess(app, None, None, task_cb, None, False)
worker.counter = 0
channel = MagicMock(callbacks={1: None, 2: None})
body = '{"task_name": "func", "args": [1], "kwargs": {"two": 2}}'
msg = MagicMock(body=body, correlation_id='testid')
self.assertRaises(SystemExit, worker._on_message, channel, msg)
task_cb.assert_called_with('func', [1], {'two': 2})
self.assertFalse(app.tasks.func.apply.called)
channel.basic_reject.assert_called_with(ANY, requeue=True)
def test_get_worker_app(self, exit_mock, fork_mock):
app = MagicMock()
worker = _WorkerProcess(app, ['testqueue'], None, None, None, False)
worker._run = MagicMock()
master = WorkerMaster('testapp')
fork_mock.return_value = 0
self.assertRaises(RuntimeError, get_worker_app)
master._start_worker(worker)
self.assertEquals(app, get_worker_app())
fork_mock.assert_called_with()
worker._run.assert_called_with()
exit_mock.assert_called_with(os.EX_OK)
def test_on_message_task_fail(self):
app = MagicMock()
return_cb = MagicMock()
app.tasks.func.apply.side_effect = ValueError
worker = _WorkerProcess(app, None, 2, None, return_cb, False)
worker.counter = 0
channel = MagicMock()
body = '{"task_name": "func", "args": [1], "kwargs": {"two": 2}}'
msg = MagicMock(body=body, correlation_id='testid')
self.assertRaises(ValueError, worker._on_message, channel, msg)
app.tasks.func.apply.assert_called_with([1], {'two': 2}, 'testid')
return_cb.assert_called_with('func', None)
channel.basic_reject.assert_called_with(ANY, requeue=False)
self.assertFalse(channel.basic_cancel.called)
def test_on_message_task_callback_discardtask(self):
app = MagicMock()
task_cb = MagicMock(side_effect=DiscardTask)
return_cb = MagicMock()
worker = _WorkerProcess(app, None, None, task_cb, return_cb,
False)
worker.counter = 0
channel = MagicMock(callbacks={1: None, 2: None})
body = '{"task_name": "func", "args": [1], "kwargs": {"two": 2}}'
msg = MagicMock(body=body, correlation_id='testid')
worker._on_message(channel, msg)
task_cb.assert_called_with('func', [1], {'two': 2})
self.assertFalse(app.tasks.func.apply.called)
self.assertFalse(return_cb.called)
channel.basic_ack.assert_called_with(ANY)
def test_on_message_limit(self):
app = MagicMock()
worker = _WorkerProcess(app, None, 1, None, None, False)
worker.counter = 0
channel = MagicMock(callbacks={1: None, 2: None})
def clear_callbacks(consumer_tag):
channel.callbacks = {}
channel.basic_cancel.side_effect = clear_callbacks
body = '{"task_name": "func", "args": [1], "kwargs": {"two": 2}}'
msg = MagicMock(body=body, correlation_id='testid')
worker._on_message(channel, msg)
app.tasks.func.apply.assert_called_with([1], {'two': 2}, 'testid')
channel.basic_ack.assert_called_with(ANY)
channel.basic_cancel.assert_any_call(1)
channel.basic_cancel.assert_any_call(2)
def test_on_message(self):
app = MagicMock()
task_cb = MagicMock()
return_cb = MagicMock()
app.tasks.func.apply.return_value = 'return'
worker = _WorkerProcess(app, None, None, task_cb, return_cb, False)
worker.counter = 0
channel = MagicMock()
body = '{"task_name": "func", "args": [1], "kwargs": {"two": 2}}'
msg = MagicMock(body=body, correlation_id='testid')
worker._on_message(channel, msg)
task_cb.assert_called_with('func', [1], {'two': 2})
app.tasks.func.apply.assert_called_with([1], {'two': 2}, 'testid')
channel.basic_ack.assert_called_with(ANY)
return_cb.assert_called_with('func', 'return')
self.assertFalse(channel.basic_cancel.called)
def test_consume(self, amqp_exit_mock, amqp_enter_mock):
app = MagicMock()
worker = _WorkerProcess(app, ['testqueue'], 1, None, None, 'exclusive')
channel = MagicMock(callbacks=True)
def set_done(timeout):
channel.callbacks = False
raise socket_timeout
channel.connection.drain_events.side_effect = set_done
amqp_enter_mock.return_value = channel
worker._consume()
amqp_enter_mock.assert_called_with()
channel.basic_consume.assert_called_with(queue='testqueue',
consumer_tag='testqueue',
callback=ANY,
exclusive='exclusive')
channel.connection.drain_events.assert_called_with(timeout=10.0)
channel.connection.send_heartbeat.assert_called_with()
amqp_exit_mock.assert_called_with(None, None, None)