def test_requeue(self, ):
worker = sergeant.worker.Worker()
worker.config = sergeant.config.WorkerConfig(
name='some_worker',
connector=sergeant.config.Connector(
type='redis',
params={
'nodes': [
{
'host': 'localhost',
'port': 6379,
'password': None,
'database': 0,
},
],
},
),
max_retries=3,
)
worker.init_broker()
task = sergeant.objects.Task()
worker.broker.push_task = unittest.mock.MagicMock()
with self.assertRaises(
expected_exception=sergeant.worker.WorkerRequeue, ):
worker.requeue(task=task, )
self.assertEqual(
first=worker.broker.push_task.call_count,
second=1,
)
def test_actions(self, ):
worker = sergeant.worker.Worker()
worker.config = sergeant.config.WorkerConfig(
name='some_worker',
connector=sergeant.config.Connector(
type='redis',
params={
'nodes': [
{
'host': 'localhost',
'port': 6379,
'password': None,
'database': 0,
},
],
},
),
)
worker.init_broker()
worker.purge_tasks()
self.assertEqual(
first=worker.number_of_enqueued_tasks(),
second=0,
)
worker.push_task(kwargs={
'task': 1,
}, )
self.assertEqual(
first=worker.number_of_enqueued_tasks(),
second=1,
)
worker.push_tasks(kwargs_list=[
{
'task': 2,
},
{
'task': 3,
},
{
'task': 4,
},
], )
self.assertEqual(
first=worker.number_of_enqueued_tasks(),
second=4,
)
tasks = list(worker.get_next_tasks(number_of_tasks=1, ))
self.assertEqual(
first=tasks[0].kwargs['task'],
second=1,
)
worker.purge_tasks()
self.assertEqual(
first=worker.number_of_enqueued_tasks(),
second=0,
)
worker.purge_tasks(task_name='other_worker', )
self.assertEqual(
first=worker.number_of_enqueued_tasks(task_name='other_worker', ),
second=0,
)
worker.push_task(
task_name='other_worker',
kwargs={
'task': 1,
},
)
self.assertEqual(
first=worker.number_of_enqueued_tasks(task_name='other_worker', ),
second=1,
)
worker.push_tasks(
task_name='other_worker',
kwargs_list=[
{
'task': 2,
},
{
'task': 3,
},
{
'task': 4,
},
],
)
self.assertEqual(
first=worker.number_of_enqueued_tasks(task_name='other_worker', ),
second=4,
)
tasks = list(
worker.get_next_tasks(
task_name='other_worker',
number_of_tasks=1,
))
self.assertEqual(
first=tasks[0].kwargs['task'],
second=1,
)
worker.purge_tasks(task_name='other_worker', )
self.assertEqual(
first=worker.number_of_enqueued_tasks(task_name='other_worker', ),
second=0,
)
lock = worker.lock(name='test_lock', )
lock.release()
self.assertFalse(expr=lock.is_locked(), )
self.assertTrue(expr=lock.acquire(timeout=0, ), )
self.assertFalse(expr=lock.acquire(timeout=0, ), )
self.assertTrue(expr=lock.release(), )
self.assertFalse(expr=lock.release(), )
self.assertIsNone(obj=lock.get_ttl(), )
self.assertTrue(expr=lock.acquire(
timeout=0,
ttl=60,
), )
self.assertEqual(
first=lock.get_ttl(),
second=60,
)
self.assertTrue(expr=lock.set_ttl(ttl=30, ), )
self.assertEqual(
first=lock.get_ttl(),
second=30,
)
def test_init_broker(self, ):
worker = sergeant.worker.Worker()
worker.config = sergeant.config.WorkerConfig(
name='some_worker',
connector=sergeant.config.Connector(
type='redis',
params={
'nodes': [
{
'host': 'localhost',
'port': 6379,
'password': None,
'database': 0,
},
],
},
),
)
worker.init_broker()
self.assertIsInstance(
obj=worker.broker,
cls=sergeant.broker.Broker,
)
self.assertIsInstance(
obj=worker.broker.connector,
cls=sergeant.connector.redis.Connector,
)
self.assertIsNone(obj=worker.broker.encoder.compressor, )
self.assertIsInstance(
obj=worker.broker.encoder.serializer,
cls=sergeant.encoder.serializer.pickle.Serializer,
)
worker.config = sergeant.config.WorkerConfig(
name='some_worker',
connector=sergeant.config.Connector(
type='redis',
params={
'nodes': [
{
'host': 'localhost',
'port': 6379,
'password': None,
'database': 0,
},
{
'host': 'localhost',
'port': 6380,
'password': None,
'database': 0,
},
],
},
),
)
worker.init_broker()
self.assertIsInstance(
obj=worker.broker,
cls=sergeant.broker.Broker,
)
self.assertIsInstance(
obj=worker.broker.connector,
cls=sergeant.connector.redis.Connector,
)
self.assertIsNone(obj=worker.broker.encoder.compressor, )
self.assertIsInstance(
obj=worker.broker.encoder.serializer,
cls=sergeant.encoder.serializer.pickle.Serializer,
)
worker.config = sergeant.config.WorkerConfig(
name='some_worker',
connector=sergeant.config.Connector(
type='mongo',
params={
'mongodb_uri': 'mongodb://localhost:27017/',
},
),
)
worker.init_broker()
self.assertIsInstance(
obj=worker.broker,
cls=sergeant.broker.Broker,
)
self.assertIsInstance(
obj=worker.broker.connector,
cls=sergeant.connector.mongo.Connector,
)
self.assertIsNone(obj=worker.broker.encoder.compressor, )
self.assertIsInstance(
obj=worker.broker.encoder.serializer,
cls=sergeant.encoder.serializer.pickle.Serializer,
)
compressor_names = list(
sergeant.encoder.encoder.Encoder.compressors.keys())
compressor_names.append(None)
serializer_names = sergeant.encoder.encoder.Encoder.serializers.keys()
for compressor_name in compressor_names:
for serializer_name in serializer_names:
worker.config = sergeant.config.WorkerConfig(
name='some_worker',
connector=sergeant.config.Connector(
type='redis',
params={
'nodes': [
{
'host': 'localhost',
'port': 6379,
'password': None,
'database': 0,
},
],
},
),
encoder=sergeant.config.Encoder(
compressor=compressor_name,
serializer=serializer_name,
),
)
worker.init_broker()
self.assertIsInstance(
obj=worker.broker,
cls=sergeant.broker.Broker,
)
self.assertIsInstance(
obj=worker.broker.connector,
cls=sergeant.connector.redis.Connector,
)
if compressor_name:
self.assertEqual(
first=worker.broker.encoder.compressor.name,
second=compressor_name,
)
else:
self.assertIsNone(obj=worker.broker.encoder.compressor, )
self.assertEqual(
first=worker.broker.encoder.serializer.name,
second=serializer_name,
)
def test_on_success(
self,
):
worker = sergeant.worker.Worker()
worker.config = sergeant.config.WorkerConfig(
name='some_worker',
connector=sergeant.config.Connector(
type='redis',
params={
'nodes': [
{
'host': 'localhost',
'port': 6379,
'password': None,
'database': 0,
},
],
},
),
max_retries=3,
logging=sergeant.config.Logging(
events=sergeant.config.LoggingEvents(
on_success=False,
),
),
)
worker.init_broker()
task = sergeant.objects.Task()
worker.on_success = unittest.mock.MagicMock()
worker.logger = unittest.mock.MagicMock()
worker.handle_success(
task=task,
returned_value=True,
)
worker.on_success.assert_called_once()
worker.logger.error.assert_not_called()
worker.on_success.reset_mock()
worker.logger.reset_mock()
worker.config = worker.config.replace(
logging=sergeant.config.Logging(
events=sergeant.config.LoggingEvents(
on_success=True,
),
),
)
worker.handle_success(
task=task,
returned_value=True,
)
worker.on_success.assert_called_once()
worker.logger.info.assert_called_once_with(
msg='task has finished successfully',
extra={
'task': task,
},
)
worker.on_success.reset_mock()
worker.logger.reset_mock()
worker.on_success.side_effect = Exception('exception message')
worker.handle_success(
task=task,
returned_value=True,
)
worker.on_success.assert_called_once()
worker.logger.info.assert_any_call(
msg='task has finished successfully',
extra={
'task': task,
},
)
worker.logger.error.assert_any_call(
msg='on_success handler has failed: exception message',
extra={
'task': task,
},
)
worker.on_success.reset_mock()
worker.logger.reset_mock()
worker.on_success.side_effect = sergeant.worker.WorkerStop()
with self.assertRaises(
expected_exception=sergeant.worker.WorkerStop,
):
worker.handle_success(
task=task,
returned_value=None,
)
worker.on_success.assert_called_once()
worker.logger.info.assert_called_once_with(
msg='task has finished successfully',
extra={
'task': task,
},
)
worker.logger.error.assert_not_called()
worker.on_success.reset_mock()
worker.logger.reset_mock()
worker.on_success.side_effect = sergeant.worker.WorkerRespawn()
with self.assertRaises(
expected_exception=sergeant.worker.WorkerRespawn,
):
worker.handle_success(
task=task,
returned_value=None,
)
worker.on_success.assert_called_once()
worker.logger.info.assert_called_once_with(
msg='task has finished successfully',
extra={
'task': task,
},
)
worker.logger.error.assert_not_called()
def test_actions_from_handlers(
self,
):
worker = sergeant.worker.Worker()
worker.config = sergeant.config.WorkerConfig(
name='some_worker',
connector=sergeant.config.Connector(
type='redis',
params={
'nodes': [
{
'host': 'localhost',
'port': 6379,
'password': None,
'database': 0,
},
],
},
),
max_retries=3,
)
worker.init_broker()
task = sergeant.objects.Task()
worker.logger = unittest.mock.MagicMock()
worker.on_timeout = unittest.mock.MagicMock()
worker.on_failure = unittest.mock.MagicMock()
worker.handle_requeue = unittest.mock.MagicMock()
worker.on_timeout.side_effect = sergeant.worker.WorkerRequeue()
worker.handle_timeout(
task=task,
)
worker.on_timeout.assert_called_once()
worker.handle_requeue.assert_called_once()
worker.handle_requeue.reset_mock()
worker.handle_requeue = unittest.mock.MagicMock()
worker.on_failure.side_effect = sergeant.worker.WorkerRequeue()
worker.handle_failure(
task=task,
exception=Exception('test_exception'),
)
worker.on_failure.assert_called_once()
worker.handle_requeue.assert_called_once()
worker.on_timeout.reset_mock()
worker.on_failure.reset_mock()
worker.handle_retry = unittest.mock.MagicMock()
worker.on_timeout.side_effect = sergeant.worker.WorkerRetry()
worker.handle_timeout(
task=task,
)
worker.on_timeout.assert_called_once()
worker.handle_retry.assert_called_once()
worker.handle_retry.reset_mock()
worker.handle_retry = unittest.mock.MagicMock()
worker.on_failure.side_effect = sergeant.worker.WorkerRetry()
worker.handle_failure(
task=task,
exception=Exception('test_exception'),
)
worker.on_failure.assert_called_once()
worker.handle_retry.assert_called_once()
worker.on_timeout.reset_mock()
worker.on_failure.reset_mock()
worker.handle_max_retries = unittest.mock.MagicMock()
worker.on_timeout.side_effect = sergeant.worker.WorkerMaxRetries()
worker.handle_timeout(
task=task,
)
worker.on_timeout.assert_called_once()
worker.handle_max_retries.assert_called_once()
worker.handle_max_retries.reset_mock()
worker.handle_max_retries = unittest.mock.MagicMock()
worker.on_failure.side_effect = sergeant.worker.WorkerMaxRetries()
worker.handle_failure(
task=task,
exception=Exception('test_exception'),
)
worker.on_failure.assert_called_once()
worker.handle_max_retries.assert_called_once()
def test_retry(self, ):
worker = sergeant.worker.Worker()
worker.config = sergeant.config.WorkerConfig(
name='some_worker',
connector=sergeant.config.Connector(
type='redis',
params={
'nodes': [
{
'host': 'localhost',
'port': 6379,
'password': None,
'database': 0,
},
],
},
),
max_retries=3,
)
worker.init_broker()
task = sergeant.objects.Task()
worker.broker.apply_async_one = unittest.mock.MagicMock()
with self.assertRaises(
expected_exception=sergeant.worker.WorkerRetry, ):
worker.retry(task=task, )
self.assertEqual(
first=worker.broker.apply_async_one.call_count,
second=1,
)
with self.assertRaises(
expected_exception=sergeant.worker.WorkerRetry, ):
worker.retry(
task=worker.broker.apply_async_one.call_args[1]['task'], )
self.assertEqual(
first=worker.broker.apply_async_one.call_count,
second=2,
)
with self.assertRaises(
expected_exception=sergeant.worker.WorkerRetry, ):
worker.retry(
task=worker.broker.apply_async_one.call_args[1]['task'], )
self.assertEqual(
first=worker.broker.apply_async_one.call_count,
second=3,
)
with self.assertRaises(
expected_exception=sergeant.worker.WorkerMaxRetries, ):
worker.retry(
task=worker.broker.apply_async_one.call_args[1]['task'], )
self.assertEqual(
first=worker.broker.apply_async_one.call_count,
second=3,
)
worker.config = worker.config.replace(max_retries=0, )
for i in range(100):
with self.assertRaises(
expected_exception=sergeant.worker.WorkerRetry, ):
worker.retry(
task=worker.broker.apply_async_one.call_args[1]['task'], )
self.assertEqual(
first=worker.broker.apply_async_one.call_count,
second=103,
)
