def setUp(self):
self.loop = ZMQEventLoop()
asyncio.set_event_loop(self.loop)
self.worker = WorkerManagement()
self.worker.loop = self.loop
self.worker.manager = PluginManager()
class TestWorker(TestCase):
@patch('alamo_worker.manage.WorkerManagement.__init__',
MagicMock(return_value=None))
def setUp(self):
self.loop = ZMQEventLoop()
asyncio.set_event_loop(self.loop)
self.worker = WorkerManagement()
self.worker.loop = self.loop
self.worker.manager = PluginManager()
def tearDown(self):
self.loop.close()
asyncio.set_event_loop(None)
@asyncio.coroutine
def mark_soft_exit(self):
self.worker.soft_exit = True
def get_future(self, result):
"""Wrap `result` as Future object."""
fut = asyncio.Future(loop=self.loop)
fut.set_result(result)
return fut
@patch('alamo_worker.manage.WorkerManagement.__init__',
MagicMock(return_value=None))
@patch('alamo_worker.manage.WorkerManagement.close')
def test_worker_manager_as_context_manager(self, close):
with WorkerManagement():
pass
self.assertTrue(close.called)
@patch('alamo_worker.manage.PluginManager.load')
def test_load_plugins_with_correct_plugin_list(self, load_mock):
self.worker._load_plugins()
self.assertTrue(load_mock.called)
@patch('alamo_worker.manage.Queue')
@patch('alamo_worker.manage.Sentinel.master_for')
def test_connect_to_queue(self, zmq_mock, master_for):
self.worker._connect_to_queue()
self.assertTrue(master_for.called)
self.assertTrue(zmq_mock.called)
@patch('alamo_worker.manage.WorkerManagement.__init__', Mock())
@patch('alamo_worker.manage.Session.get')
def test_get_environments(self, session_mock):
mock_response = MagicMock(status_code=200)
mock_response.raise_for_status.return_value = None
mock_response.json.return_value = ENV_CONF
session_mock.return_value = mock_response
envs = self.worker._get_environments()
self.assertEqual(envs, ENV_CONF)
def test_send_method(self):
self.worker.queue = MagicMock()
self.worker.send('foo')
self.worker.queue.enqueue.assert_called_with(
'alamo_alerter.job.analyze',
args=('foo',),
result_ttl=60
)
@patch('alamo_worker.manage.SchedulerConnection')
def test_worker_open_method(self, mocked):
mocked.return_value = MagicMock()
self.worker.scheduler_pool = []
self.worker.open()
self.assertEqual(len(self.worker.scheduler_pool), 1)
def test_worker_close_method(self):
scheduler = MagicMock()
self.worker.scheduler_pool = [scheduler]
self.worker.close()
scheduler.close.assert_called_once_with()
@unpack
@data((False, 1), (True, 2))
def test_worker_generate_futures_method(self, monitor, expected):
scheduler = MagicMock()
self.worker.scheduler_pool = [scheduler]
with override_settings(SCHEDULER_MONITOR=monitor):
futures = [f for f in self.worker.generate_futures()]
self.assertEqual(len(futures), expected)
def test_worker_wait_and_kill_method(self):
self.loop.run_until_complete(self.worker.wait_and_kill('SIGINT'))
self.assertTrue(self.worker.soft_exit)
@unpack
@data(
(None, None, False, False),
({'id': 999, 'type': 'graphite'}, dict(id=999), True, True),
({'id': 999, 'type': 'graphite'}, None, False, True),
)
@patch('alamo_worker.manage.PluginManager.dispatch')
@patch('alamo_worker.manage.WorkerManagement.send')
def test_worker_receiver_method(self,
event, dispatch_res,
send_called, dispatch_called,
send, dispatch):
@asyncio.coroutine
def _dispatch(*args, **kwargs):
return dispatch_res
dispatch.side_effect = _dispatch
send._is_coroutine = False
scheduler = MagicMock()
scheduler.receive.side_effect = [
self.get_future(event),
self.mark_soft_exit()
]
self.worker.scheduler_pool = [scheduler]
future = asyncio.async(self.worker.receiver(scheduler))
self.loop.run_until_complete(future)
self.assertTrue(future.done())
self.assertEqual(send.called, send_called)
self.assertEqual(dispatch.called, dispatch_called)
@unpack
@data(
(True, dict(event=1), False, True),
(True, dict(event=EVENT_DISCONNECTED), True, True),
(True, dict(event=EVENT_MONITOR_STOPPED), False, True),
(False, dict(event=1), False, False)
)
@patch('alamo_worker.manage.logger')
def test_worker_monitor_method(self, monitor, event, warn, info, logger):
scheduler = MagicMock()
scheduler.receive_event.side_effect = [
self.get_future(event), self.mark_soft_exit()
]
self.worker.scheduler_pool = [scheduler]
future = asyncio.async(self.worker.monitor(scheduler))
with override_settings(SCHEDULER_MONITOR=monitor):
self.loop.run_until_complete(future)
self.assertTrue(future.done())
self.assertEqual(logger.warning.called, warn)
self.assertEqual(logger.info.called, info)
@patch('alamo_worker.manage.WorkerManagement.generate_futures')
def test_worker_execute_method(self, futures):
futures.return_value = [self.mark_soft_exit()]
self.worker.execute()
self.assertTrue(self.loop.is_closed())
