diff --git a/availability/storage.py b/availability/storage.py index 7e24e55..e7769a9 100644 --- a/availability/storage.py +++ b/availability/storage.py @@ -25,12 +25,16 @@ NUMBER_OF_SHARDS = 2 +class StorageException(Exception): + pass + + def get_elasticsearch(check_availability=False): """Return Elasticsearch instance. :param check_availability: check if nodes are available - :returns: Elasticsearch or None on failure - :rtype: elasticsearch.Elasticsearch + :returns: elasticsearch.Elasticsearch + :raises: StorageException """ nodes = config.get_config()["backend"]["connection"] try: @@ -38,10 +42,10 @@ def get_elasticsearch(check_availability=False): if check_availability: es.info() except Exception as e: - LOG.warning( - "Failed to query Elasticsearch nodes %s: %s" - % (nodes, str(e))) - raise + mesg = ("Failed to query Elasticsearch nodes %s: %s" + % (nodes, str(e))) + LOG.error(mesg) + raise StorageException(mesg) return es diff --git a/availability/watcher.py b/availability/watcher.py index 408956e..36a0ead 100644 --- a/availability/watcher.py +++ b/availability/watcher.py @@ -146,8 +146,9 @@ def main(period=None): LOG.error("Unexpected backend: %(type)s" % backend) return 1 - if not storage.get_elasticsearch(check_availability=True): - LOG.error("Failed to set up Elasticsearch") + try: + storage.get_elasticsearch(check_availability=True) + except storage.StorageException: return 1 LOG.info("Start watching with period %s seconds" % period) diff --git a/tests/unit/test_storage.py b/tests/unit/test_storage.py index ce16765..953ace7 100644 --- a/tests/unit/test_storage.py +++ b/tests/unit/test_storage.py @@ -62,16 +62,15 @@ def test_ensure_es_index_exists( @mock.patch("availability.storage.LOG") def test_get_elasticsearch(self, mock_log, mock_config, mock_elastic): mock_es = mock.Mock() - mock_es.indices.exists.side_effect = ValueError mock_elastic.return_value = mock_es mock_config.get_config.return_value = ( {"backend": {"connection": "nodes"}}) - self.assertFalse(mock_es.info.called) self.assertEqual(mock_es, storage.get_elasticsearch()) + self.assertFalse(mock_es.info.called) - mock_es.indices.exists.side_effect = None mock_elastic.reset_mock() - self.assertEqual(mock_es, storage.get_elasticsearch( - check_availability=True)) + mock_es.info.side_effect = ValueError + self.assertRaises(storage.StorageException, + storage.get_elasticsearch, check_availability=True) mock_elastic.assert_called_once_with("nodes") mock_es.info.assert_called_once_with() diff --git a/tests/unit/test_watcher.py b/tests/unit/test_watcher.py index ade123e..8305136 100644 --- a/tests/unit/test_watcher.py +++ b/tests/unit/test_watcher.py @@ -16,6 +16,7 @@ import mock import requests +from availability import storage from availability import watcher from tests.unit import test @@ -134,13 +135,13 @@ def test_watch_services(self, mock_log, mock_queue, mock_thread, self.assertEqual(calls, mock_thread.mock_calls) @mock.patch("availability.watcher.schedule") - @mock.patch("availability.watcher.storage") + @mock.patch("availability.watcher.storage.get_elasticsearch") @mock.patch("availability.watcher.watch_services") @mock.patch("availability.watcher.config.get_config") @mock.patch("availability.watcher.time") @mock.patch("availability.watcher.LOG") def test_main(self, mock_log, mock_time, mock_get_config, - mock_watch_services, mock_storage, mock_schedule): + mock_watch_services, mock_get_elastic, mock_schedule): class BreakInfinityCicle(Exception): pass @@ -166,19 +167,18 @@ class BreakInfinityCicle(Exception): backend = {"type": "elastic", "connection": "foo_conn"} mock_get_config.return_value = {"regions": ["foo_region"], "backend": backend} - mock_storage.get_elasticsearch.return_value = None + mock_get_elastic.side_effect = storage.StorageException self.assertEqual(1, watcher.main()) - mock_storage.get_elasticsearch.assert_called_once_with( + mock_get_elastic.assert_called_once_with( check_availability=True) - mock_storage.reset_mock() - mock_storage.get_elasticsearch.return_value = mock.Mock() + mock_get_elastic.reset_mock() + mock_get_elastic.side_effect = None self.assertRaises(BreakInfinityCicle, watcher.main) self.assertEqual([mock.call(1)] * 4, mock_time.sleep.mock_calls) mock_schedule.every.assert_called_once_with(60) - mock_storage.get_elasticsearch.assert_called_once_with( - check_availability=True) + mock_get_elastic.assert_called_once_with(check_availability=True) mock_get_config.return_value = {"regions": ["foo_region"], "backend": backend, "period": 42}