diff --git a/datahub/core/management/commands/rq_health_check.py b/datahub/core/management/commands/rq_health_check.py new file mode 100644 index 000000000..456fa6d35 --- /dev/null +++ b/datahub/core/management/commands/rq_health_check.py @@ -0,0 +1,42 @@ +import sys + +from functools import reduce +from logging import getLogger +from operator import concat + +from django.conf import settings +from django.core.management.base import BaseCommand +from redis import Redis +from rq import Worker + + +logger = getLogger(__name__) + + +class Command(BaseCommand): + help = 'RQ Health Check' + + def add_arguments(self, parser): + """Define extra arguments.""" + parser.add_argument( + '--queue', + type=str, + help='Name of the queue to perform health check on.', + ) + + def handle(self, *args, **options): + if options['queue']: + queue = str(options['queue']) + redis = Redis.from_url(settings.REDIS_BASE_URL) + workers = Worker.all(connection=redis) + queue_names = reduce(concat, [worker.queue_names() for worker in workers], []) + missing_queues = set([queue]) - set(queue_names) + + if missing_queues: + logger.error(f'RQ queue not running: {missing_queues}') + sys.exit(1) + logger.info('OK') + sys.exit(0) + + logger.error('Nothing checked! Please provide --queue parameter') + sys.exit(1) diff --git a/datahub/core/test/management/commands/test_rq_health_check.py b/datahub/core/test/management/commands/test_rq_health_check.py new file mode 100644 index 000000000..e42e2e631 --- /dev/null +++ b/datahub/core/test/management/commands/test_rq_health_check.py @@ -0,0 +1,63 @@ +import logging +from unittest import mock +from unittest.mock import patch + +import pytest + +from django.core.management import call_command + + +class MockWorker: + """ + Mock queue names object returned by worker + """ + + queue_name = '' + + def __init__(self, queue_name, *args, **kwargs): + self.queue_name = queue_name + + def queue_names(self): + return self.queue_name + + +def test_rq_health_check_ok(): + logger = logging.getLogger('datahub.core.management.commands.rq_health_check') + with patch( + 'datahub.core.management.commands.rq_health_check.Worker.all', + return_value=[MockWorker(['short-running']), MockWorker(['long-running'])], + ): + with mock.patch.object(logger, 'info') as mock_info: + with pytest.raises(SystemExit) as exception_info: + call_command('rq_health_check', '--queue=short-running') + + assert exception_info.value.code == 0 + assert 'OK' in str(mock_info.call_args_list) + assert mock_info.call_count == 1 + + +def test_rq_health_check_rq_not_running(): + logger = logging.getLogger('datahub.core.management.commands.rq_health_check') + with patch( + 'datahub.core.management.commands.rq_health_check.Worker.all', + return_value=[MockWorker(['long-running'])], + ): + with mock.patch.object(logger, 'error') as mock_error: + with pytest.raises(SystemExit) as exception_info: + call_command('rq_health_check', '--queue=short-running') + + assert exception_info.value.code == 1 + assert "RQ queue not running: {\'short-running\'}" in str(mock_error.call_args_list) + assert mock_error.call_count == 1 + + +def test_command_called_without_parameter(): + logger = logging.getLogger('datahub.core.management.commands.rq_health_check') + with mock.patch.object(logger, 'error') as mock_error: + with pytest.raises(SystemExit) as exception_info: + call_command('rq_health_check') + + assert exception_info.value.code == 1 + assert 'Nothing checked! Please provide --queue parameter' \ + in str(mock_error.call_args_list) + assert mock_error.call_count == 1 diff --git a/datahub/ping/test/test_ping_view.py b/datahub/ping/test/test_ping_view.py index 9bd529c14..f725620de 100644 --- a/datahub/ping/test/test_ping_view.py +++ b/datahub/ping/test/test_ping_view.py @@ -1,5 +1,9 @@ +from unittest.mock import patch + import pytest +from django.db import DatabaseError + from rest_framework import status from rest_framework.reverse import reverse @@ -10,6 +14,20 @@ def test_all_good(client): """Test all good.""" url = reverse('ping') response = client.get(url) + assert response.status_code == status.HTTP_200_OK assert 'OK' in str(response.content) assert response.headers['content-type'] == 'text/xml' + + +def test_check_database_fail(client): + url = reverse('ping') + with patch( + 'datahub.ping.services.Company.objects.all', + side_effect=DatabaseError('No database'), + ): + response = client.get(url) + + assert response.status_code == status.HTTP_500_INTERNAL_SERVER_ERROR + assert 'FALSE' in str(response.content) + assert response.headers['content-type'] == 'text/xml' diff --git a/docker-compose.yml b/docker-compose.yml index b0fd1f3a8..165d1e76b 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -26,6 +26,16 @@ services: depends_on: - api command: python short-running-worker.py + healthcheck: + test: + [ + "CMD-SHELL", + "python ./manage.py rq_health_check --queue=short-running" + ] + interval: 10s + timeout: 5s + retries: 2 + start_period: 5s rq_long: build: @@ -37,6 +47,16 @@ services: depends_on: - api command: python long-running-worker.py + healthcheck: + test: + [ + "CMD-SHELL", + "python ./manage.py rq_health_check --queue=long-running" + ] + interval: 10s + timeout: 5s + retries: 2 + start_period: 5s rq_sched: build: