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())