Add RabbitMQ health check (#197)

README.rst

@@ -15,6 +15,7 @@ The following health checks are bundled with this project:
 - disk and memory utilization (via ``psutil``)
 - AWS S3 storage
 - Celery task queue
+- RabbitMQ
 
 Writing your own custom health checks is also very quick and easy.
 
@@ -71,6 +72,7 @@ Add the ``health_check`` applications to your ``INSTALLED_APPS``:
         'health_check.contrib.celery',              # requires celery
         'health_check.contrib.psutil',              # disk and memory utilization; requires psutil
         'health_check.contrib.s3boto_storage',      # requires boto and S3BotoStorage backend
+        'health_check.contrib.rabbitmq',            # requires RabbitMQ broker
     ]
 
 (Optional) If using the ``psutil`` app, you can configure disk and memory
@@ -90,6 +92,14 @@ If using the DB check, run migrations:
 
     django-admin migrate
 
+To use the RabbitMQ healthcheck, please make sure that there is a variable named ``BROKER_URL``
+on django.conf.settings with the required format to connect to your rabbit server. For example:
+
+.. code::
+
+    BROKER_URL = amqp://myuser:mypassword@localhost:5672/myvhost
+
+
 Setting up monitoring
 ---------------------
 

health_check/contrib/rabbitmq/__init__.py

@@ -0,0 +1 @@
+default_app_config = 'health_check.contrib.rabbitmq.apps.HealthCheckConfig'

health_check/contrib/rabbitmq/apps.py

@@ -0,0 +1,12 @@
+from django.apps import AppConfig
+
+from health_check.plugins import plugin_dir
+
+
+class HealthCheckConfig(AppConfig):
+    name = "health_check.contrib.rabbitmq"
+
+    def ready(self):
+        from .backends import RabbitMQHealthCheck
+
+        plugin_dir.register(RabbitMQHealthCheck)

health_check/contrib/rabbitmq/backends.py

@@ -0,0 +1,41 @@
+import logging
+
+from amqp.exceptions import AccessRefused
+from django.conf import settings
+from kombu import Connection
+
+from health_check.backends import BaseHealthCheckBackend
+from health_check.exceptions import ServiceUnavailable
+
+logger = logging.getLogger(__name__)
+
+
+class RabbitMQHealthCheck(BaseHealthCheckBackend):
+    """Health check for RabbitMQ."""
+
+    def check_status(self):
+        """Check RabbitMQ service by opening and closing a broker channel."""
+        logger.debug("Checking for a broker_url on django settings...")
+
+        broker_url = getattr(settings, "BROKER_URL", None)
+
+        logger.debug("Got %s as the broker_url. Connecting to rabbit...", broker_url)
+
+        logger.debug("Attempting to connect to rabbit...")
+        try:
+            # conn is used as a context to release opened resources later
+            with Connection(broker_url) as conn:
+                conn.connect()  # exceptions may be raised upon calling connect
+        except ConnectionRefusedError as e:
+            self.add_error(ServiceUnavailable("Unable to connect to RabbitMQ: Connection was refused."), e)
+
+        except AccessRefused as e:
+            self.add_error(ServiceUnavailable("Unable to connect to RabbitMQ: Authentication error."), e)
+
+        except IOError as e:
+            self.add_error(ServiceUnavailable("IOError"), e)
+
+        except BaseException as e:
+            self.add_error(ServiceUnavailable("Unknown error"), e)
+        else:
+            logger.debug("Connection estabilished. RabbitMQ is healthy.")

tests/test_rabbitmq.py

@@ -0,0 +1,77 @@
+import mock
+from amqp.exceptions import AccessRefused
+
+from health_check.contrib.rabbitmq.backends import RabbitMQHealthCheck
+
+
+class TestRabbitMQHealthCheck:
+    """Test RabbitMQ health check."""
+
+    @mock.patch("health_check.contrib.rabbitmq.backends.getattr")
+    @mock.patch("health_check.contrib.rabbitmq.backends.Connection")
+    def test_broker_refused_connection(self, mocked_connection, mocked_getattr):
+        """Test when the connection to RabbitMQ is refused."""
+        mocked_getattr.return_value = "broker_url"
+
+        conn_exception = ConnectionRefusedError("Refused connection")
+
+        # mock returns
+        mocked_conn = mock.MagicMock()
+        mocked_connection.return_value.__enter__.return_value = mocked_conn
+        mocked_conn.connect.side_effect = conn_exception
+
+        # instantiates the class
+        rabbitmq_healthchecker = RabbitMQHealthCheck()
+
+        # invokes the method check_status()
+        rabbitmq_healthchecker.check_status()
+        assert len(rabbitmq_healthchecker.errors), 1
+
+        # mock assertions
+        mocked_connection.assert_called_once_with("broker_url")
+
+    @mock.patch("health_check.contrib.rabbitmq.backends.getattr")
+    @mock.patch("health_check.contrib.rabbitmq.backends.Connection")
+    def test_broker_auth_error(self, mocked_connection, mocked_getattr):
+        """Test that the connection to RabbitMQ has an authentication error."""
+        mocked_getattr.return_value = "broker_url"
+
+        conn_exception = AccessRefused("Refused connection")
+
+        # mock returns
+        mocked_conn = mock.MagicMock()
+        mocked_connection.return_value.__enter__.return_value = mocked_conn
+        mocked_conn.connect.side_effect = conn_exception
+
+        # instantiates the class
+        rabbitmq_healthchecker = RabbitMQHealthCheck()
+
+        # invokes the method check_status()
+        rabbitmq_healthchecker.check_status()
+        assert len(rabbitmq_healthchecker.errors), 1
+
+        # mock assertions
+        mocked_connection.assert_called_once_with("broker_url")
+
+    @mock.patch("health_check.contrib.rabbitmq.backends.getattr")
+    @mock.patch("health_check.contrib.rabbitmq.backends.Connection")
+    def test_broker_connection_upon_none_url(self, mocked_connection, mocked_getattr):
+        """Thest when the connection to RabbitMQ has no ``broker_url``."""
+        mocked_getattr.return_value = None
+        # if the variable BROKER_URL is not set, AccessRefused exception is raised
+        conn_exception = AccessRefused("Refused connection")
+
+        # mock returns
+        mocked_conn = mock.MagicMock()
+        mocked_connection.return_value.__enter__.return_value = mocked_conn
+        mocked_conn.connect.side_effect = conn_exception
+
+        # instantiates the class
+        rabbitmq_healthchecker = RabbitMQHealthCheck()
+
+        # invokes the method check_status()
+        rabbitmq_healthchecker.check_status()
+        assert len(rabbitmq_healthchecker.errors), 1
+
+        # mock assertions
+        mocked_connection.assert_called_once_with(None)