From 338111dcd9b60b558d070f9feae552d27fbf2681 Mon Sep 17 00:00:00 2001 From: Danang Massandy Date: Tue, 27 Jan 2026 23:29:15 +0000 Subject: [PATCH 1/3] add health readiness API endpoint --- django_project/core/settings/apps/project.py | 1 + django_project/core/urls.py | 1 + django_project/geosight/health/__init__.py | 27 ++ django_project/geosight/health/apps.py | 24 ++ .../geosight/health/test_readiness.py | 247 ++++++++++++++++++ django_project/geosight/health/urls.py | 22 ++ django_project/geosight/health/views.py | 171 ++++++++++++ 7 files changed, 493 insertions(+) create mode 100644 django_project/geosight/health/__init__.py create mode 100644 django_project/geosight/health/apps.py create mode 100644 django_project/geosight/health/test_readiness.py create mode 100644 django_project/geosight/health/urls.py create mode 100644 django_project/geosight/health/views.py diff --git a/django_project/core/settings/apps/project.py b/django_project/core/settings/apps/project.py index fab901699..16d616c98 100644 --- a/django_project/core/settings/apps/project.py +++ b/django_project/core/settings/apps/project.py @@ -24,5 +24,6 @@ 'geosight.permission', 'geosight.importer', 'geosight.log', + 'geosight.health', 'frontend' ] diff --git a/django_project/core/urls.py b/django_project/core/urls.py index 63ff590b3..6df8d6eb3 100644 --- a/django_project/core/urls.py +++ b/django_project/core/urls.py @@ -66,6 +66,7 @@ def get_schema(self, request=None, public=False): # noqa DOC101 urlpatterns = [ url(r'^i18n/', include('django.conf.urls.i18n')), + path('', include('geosight.health.urls')), ] urlpatterns += i18n_patterns( diff --git a/django_project/geosight/health/__init__.py b/django_project/geosight/health/__init__.py new file mode 100644 index 000000000..55494da04 --- /dev/null +++ b/django_project/geosight/health/__init__.py @@ -0,0 +1,27 @@ +# coding=utf-8 +""" +GeoSight is UNICEF's geospatial web-based business intelligence platform. + +Contact : geosight-no-reply@unicef.org + +.. note:: This program is free software; you can redistribute it and/or modify + it under the terms of the GNU Affero General Public License as published by + the Free Software Foundation; either version 3 of the License, or + (at your option) any later version. + +""" +__author__ = 'danang@kartoza.com' +__date__ = '22/10/2024' +__copyright__ = ('Copyright 2023, Unicef') + +from django.apps import AppConfig + + +class Config(AppConfig): + """Health check application configuration.""" + + label = 'geosight_health' + name = 'geosight.health' + verbose_name = "GeoSight Health" + +default_app_config = 'geosight.health.Config' diff --git a/django_project/geosight/health/apps.py b/django_project/geosight/health/apps.py new file mode 100644 index 000000000..63d2f104e --- /dev/null +++ b/django_project/geosight/health/apps.py @@ -0,0 +1,24 @@ +# coding=utf-8 +""" +GeoSight is UNICEF's geospatial web-based business intelligence platform. + +Contact : geosight-no-reply@unicef.org + +.. note:: This program is free software; you can redistribute it and/or modify + it under the terms of the GNU Affero General Public License as published by + the Free Software Foundation; either version 3 of the License, or + (at your option) any later version. + +""" +__author__ = 'danang@kartoza.com' +__date__ = '22/10/2024' +__copyright__ = ('Copyright 2023, Unicef') + +from django.apps import AppConfig + + +class HealthConfig(AppConfig): + """Health check application configuration.""" + + default_auto_field = 'django.db.models.BigAutoField' + name = 'health' diff --git a/django_project/geosight/health/test_readiness.py b/django_project/geosight/health/test_readiness.py new file mode 100644 index 000000000..73f661d22 --- /dev/null +++ b/django_project/geosight/health/test_readiness.py @@ -0,0 +1,247 @@ +# coding=utf-8 +""" +GeoSight is UNICEF's geospatial web-based business intelligence platform. + +Contact : geosight-no-reply@unicef.org + +.. note:: This program is free software; you can redistribute it and/or modify + it under the terms of the GNU Affero General Public License as published by + the Free Software Foundation; either version 3 of the License, or + (at your option) any later version. + +""" +__author__ = 'danang@kartoza.com' +__date__ = '22/10/2024' +__copyright__ = ('Copyright 2023, Unicef') + +from django.test import TestCase, override_settings +from django.urls import reverse +from django.core.cache import cache +from unittest.mock import patch, MagicMock +import tempfile + + +@override_settings( + CACHES={ + 'default': { + 'BACKEND': 'django.core.cache.backends.locmem.LocMemCache', + 'LOCATION': 'unique-test-cache', + } + } +) +class ReadinessProbeTest(TestCase): + """Test suite for readiness probe endpoint.""" + + def setUp(self): + """Set up test environment.""" + self.url = reverse('health-readiness') + # Clear cache before each test + cache.clear() + + def tearDown(self): + """Clean up after tests.""" + cache.clear() + + def test_readiness_probe_all_healthy(self): + """Test readiness probe when all checks pass.""" + response = self.client.get(self.url) + self.assertEqual(response.status_code, 200) + self.assertEqual(response.json()['status'], 'ready') + self.assertTrue(response.json()['checks']['database']) + self.assertTrue(response.json()['checks']['redis']) + + def test_readiness_probe_returns_json(self): + """Test that readiness probe returns valid JSON.""" + response = self.client.get(self.url) + self.assertEqual(response['Content-Type'], 'application/json') + data = response.json() + self.assertIn('status', data) + self.assertIn('checks', data) + + def test_database_check_passes(self): + """Test database connectivity check passes.""" + response = self.client.get(self.url) + self.assertEqual(response.status_code, 200) + self.assertTrue(response.json()['checks']['database']) + + @patch('health.views.connection') + def test_database_check_fails(self, mock_connection): + """Test readiness probe fails when database is down. + + :param mock_connection: Mocked database connection + :type mock_connection: MagicMock + """ + # Mock database connection failure + mock_cursor = MagicMock() + mock_cursor.execute.side_effect = Exception( + "Database connection failed" + ) + mock_connection.cursor.return_value.__enter__.return_value = ( + mock_cursor + ) + response = self.client.get(self.url) + self.assertEqual(response.status_code, 503) + self.assertEqual(response.json()['status'], 'not ready') + self.assertFalse(response.json()['checks']['database']) + + def test_redis_check_passes(self): + """Test Redis connectivity check passes.""" + response = self.client.get(self.url) + self.assertEqual(response.status_code, 200) + self.assertTrue(response.json()['checks']['redis']) + + @patch('health.views.cache') + def test_redis_check_fails_on_set(self, mock_cache): + """Test readiness probe fails when Redis set operation fails. + + :param mock_cache: Mocked cache object + :type mock_cache: MagicMock + """ + mock_cache.set.side_effect = Exception("Redis connection failed") + response = self.client.get(self.url) + self.assertEqual(response.status_code, 503) + self.assertEqual(response.json()['status'], 'not ready') + self.assertFalse(response.json()['checks']['redis']) + + @patch('health.views.cache') + def test_redis_check_fails_on_get(self, mock_cache): + """Test readiness probe fails when Redis get returns wrong value. + + :param mock_cache: Mocked cache object + :type mock_cache: MagicMock + """ + mock_cache.set.return_value = None + mock_cache.get.return_value = 'wrong_value' + response = self.client.get(self.url) + + self.assertEqual(response.status_code, 503) + self.assertEqual(response.json()['status'], 'not ready') + self.assertFalse(response.json()['checks']['redis']) + + def test_storage_check_passes_with_low_usage(self): + """Test storage check passes when disk usage is below threshold.""" + with tempfile.TemporaryDirectory() as temp_dir: + with override_settings( + LOGS_DIRECTORY=temp_dir, + STORAGE_CRITICAL_THRESHOLD=98 + ): + response = self.client.get(self.url) + self.assertEqual(response.status_code, 200) + self.assertTrue(response.json()['checks']['storage']) + + @patch('health.views.shutil.disk_usage') + def test_storage_check_fails_when_usage_exceeds_threshold( + self, mock_disk_usage + ): + """Test storage check fails when disk usage exceeds 98%. + + :param mock_disk_usage: Mocked disk usage function + :type mock_disk_usage: MagicMock + """ + # Mock disk usage at 99% (exceeds 98% threshold) + mock_disk_usage.return_value = MagicMock( + total=10 * 1024**3, # 10GB total + used=9.9 * 1024**3, # 9.9GB used (99%) + free=0.1 * 1024**3 # 0.1GB free + ) + with tempfile.TemporaryDirectory() as temp_dir: + with override_settings( + LOGS_DIRECTORY=temp_dir, + STORAGE_CRITICAL_THRESHOLD=98 + ): + response = self.client.get(self.url) + self.assertEqual(response.status_code, 503) + self.assertEqual(response.json()['status'], 'not ready') + self.assertFalse(response.json()['checks']['storage']) + + @patch('health.views.shutil.disk_usage') + def test_storage_check_passes_at_threshold_boundary(self, mock_disk_usage): + """Test storage check passes when exactly at 98% (not exceeding). + + :param mock_disk_usage: Mocked disk usage function + :type mock_disk_usage: MagicMock + """ + # Mock disk usage at exactly 98% + mock_disk_usage.return_value = MagicMock( + total=10 * 1024**3, # 10GB total + used=9.8 * 1024**3, # 9.8GB used (98%) + free=0.2 * 1024**3 # 0.2GB free + ) + response = self.client.get(self.url) + self.assertEqual(response.status_code, 200) + self.assertTrue(response.json()['checks']['storage']) + + @patch('health.views.shutil.disk_usage') + def test_storage_check_with_custom_threshold(self, mock_disk_usage): + """Test storage check respects custom threshold setting. + + :param mock_disk_usage: Mocked disk usage function + :type mock_disk_usage: MagicMock + """ + # Mock disk usage at 91% + mock_disk_usage.return_value = MagicMock( + total=10 * 1024**3, + used=9.1 * 1024**3, + free=0.9 * 1024**3 + ) + with tempfile.TemporaryDirectory() as temp_dir: + # Test with 90% threshold - should fail + with override_settings( + LOGS_DIRECTORY=temp_dir, + STORAGE_CRITICAL_THRESHOLD=90 + ): + response = self.client.get(self.url) + self.assertEqual(response.status_code, 503) + self.assertFalse(response.json()['checks']['storage']) + + @patch('health.views.connection') + @patch('health.views.cache') + def test_multiple_checks_fail(self, mock_cache, mock_connection): + """Test readiness probe when multiple checks fail. + + :param mock_cache: Mocked cache object + :type mock_cache: MagicMock + :param mock_connection: Mocked database connection + :type mock_connection: MagicMock + """ + # Mock both database and Redis failures + mock_cursor = MagicMock() + mock_cursor.execute.side_effect = Exception("Database failed") + mock_connection.cursor.return_value.__enter__.return_value = ( + mock_cursor + ) + mock_cache.set.side_effect = Exception("Redis failed") + + response = self.client.get(self.url) + self.assertEqual(response.status_code, 503) + self.assertEqual(response.json()['status'], 'not ready') + self.assertFalse(response.json()['checks']['database']) + self.assertFalse(response.json()['checks']['redis']) + + def test_endpoint_allows_unauthenticated_access(self): + """Test that readiness endpoint doesn't require authentication.""" + # Should work without any authentication + response = self.client.get(self.url) + # Should not return 401 or 403 + self.assertNotEqual(response.status_code, 401) + self.assertNotEqual(response.status_code, 403) + # Should return either 200 or 503 + self.assertIn(response.status_code, [200, 503]) + + def test_endpoint_only_accepts_get(self): + """Test that readiness endpoint only accepts GET requests.""" + # GET should work + get_response = self.client.get(self.url) + self.assertIn(get_response.status_code, [200, 503]) + + # POST should not be allowed + post_response = self.client.post(self.url) + self.assertEqual(post_response.status_code, 405) + + # PUT should not be allowed + put_response = self.client.put(self.url) + self.assertEqual(put_response.status_code, 405) + + # DELETE should not be allowed + delete_response = self.client.delete(self.url) + self.assertEqual(delete_response.status_code, 405) diff --git a/django_project/geosight/health/urls.py b/django_project/geosight/health/urls.py new file mode 100644 index 000000000..70a0295e8 --- /dev/null +++ b/django_project/geosight/health/urls.py @@ -0,0 +1,22 @@ +# coding=utf-8 +""" +GeoSight is UNICEF's geospatial web-based business intelligence platform. + +Contact : geosight-no-reply@unicef.org + +.. note:: This program is free software; you can redistribute it and/or modify + it under the terms of the GNU Affero General Public License as published by + the Free Software Foundation; either version 3 of the License, or + (at your option) any later version. + +""" +__author__ = 'danang@kartoza.com' +__date__ = '22/10/2024' +__copyright__ = ('Copyright 2023, Unicef') + +from django.urls import path +from . import views + +urlpatterns = [ + path('ready/', views.readiness_probe, name='health-readiness'), +] diff --git a/django_project/geosight/health/views.py b/django_project/geosight/health/views.py new file mode 100644 index 000000000..2d0af0c9b --- /dev/null +++ b/django_project/geosight/health/views.py @@ -0,0 +1,171 @@ +# coding=utf-8 +""" +GeoSight is UNICEF's geospatial web-based business intelligence platform. + +Contact : geosight-no-reply@unicef.org + +.. note:: This program is free software; you can redistribute it and/or modify + it under the terms of the GNU Affero General Public License as published by + the Free Software Foundation; either version 3 of the License, or + (at your option) any later version. + +""" +__author__ = 'danang@kartoza.com' +__date__ = '22/10/2024' +__copyright__ = ('Copyright 2023, Unicef') + +from rest_framework.decorators import api_view, permission_classes +from rest_framework.permissions import AllowAny +from rest_framework.response import Response +from rest_framework import status +from django.db import connection +from django.core.cache import cache +from django.conf import settings +import logging +import shutil +import os + +logger = logging.getLogger(__name__) +# Add more paths if necessary +STORAGE_PATHS = { + 'logs': getattr(settings, 'LOGS_DIRECTORY', None), + 'media': getattr(settings, 'MEDIA_ROOT', None), + 'static': getattr(settings, 'STATIC_ROOT', None), +} + + +@api_view(['GET']) +@permission_classes([AllowAny]) +def readiness_probe(request): + """Readiness probe - checks if the application is ready. + + :param request: HTTP request + :type request: Request + :return: HTTP response with readiness status + :rtype: Response + """ + checks = { + "database": check_database(), + "redis": check_redis(), + "storage": check_storage(), + } + + storage_info = get_storage_info() + + # Filter out None values (optional checks) + checks = {k: v for k, v in checks.items() if v is not None} + all_healthy = all(checks.values()) + + response_data = { + "status": "ready" if all_healthy else "not ready", + "checks": checks, + } + if storage_info: + response_data["storage_info"] = storage_info + + return Response( + response_data, + status=( + status.HTTP_200_OK if all_healthy else + status.HTTP_503_SERVICE_UNAVAILABLE + ) + ) + + +def check_database(): + """Check database connectivity. + + :return: True if database is reachable, False otherwise + :rtype: bool + """ + try: + with connection.cursor() as cursor: + cursor.execute("SELECT 1") + return True + except Exception as e: + logger.error(f"Database health check failed: {str(e)}") + return False + + +def check_redis(): + """Check Redis connectivity. + + :return: True if Redis is reachable, False otherwise + :rtype: bool + """ + try: + cache.set('health_check', 'ok', 10) + result = cache.get('health_check') + + if result != 'ok': + logger.error( + "Redis health check failed: could not retrieve test value" + ) + return False + + cache.delete('health_check') + return True + except Exception as e: + logger.error(f"Redis health check failed: {str(e)}") + return False + + +def check_storage(): + """Check storage disk space - fail if usage > 98%. + + :return: True if storage usage is below threshold, False otherwise + :rtype: bool + """ + try: + critical_threshold = getattr( + settings, 'STORAGE_CRITICAL_THRESHOLD', 98 + ) + all_healthy = True + + for name, path in STORAGE_PATHS.items(): + if not path or not os.path.exists(path): + continue + + disk_usage = shutil.disk_usage(path) + usage_percent = round( + (disk_usage.used / disk_usage.total) * 100, 2 + ) + + if usage_percent > critical_threshold: + logger.error( + f"Storage '{name}' CRITICAL: {usage_percent:.2f}% " + f"used at {path}" + ) + all_healthy = False + else: + logger.debug(f"Storage '{name}' OK: {usage_percent:.2f}% used") + + return all_healthy + except Exception as e: + logger.error(f"Storage health check failed: {str(e)}") + return False + + +def get_storage_info(): + """Get detailed info for all storage locations. + + :return: Dictionary with storage info or None if no paths available + :rtype: dict or None + """ + storage_info = {} + for name, path in STORAGE_PATHS.items(): + if not path or not os.path.exists(path): + continue + try: + disk_usage = shutil.disk_usage(path) + usage_percent = (disk_usage.used / disk_usage.total) * 100 + storage_info[name] = { + "path": path, + "total_gb": round(disk_usage.total / (1024**3), 2), + "used_gb": round(disk_usage.used / (1024**3), 2), + "free_gb": round(disk_usage.free / (1024**3), 2), + "usage_percent": round(usage_percent, 2), + } + except Exception as e: + logger.warning(f"Could not get storage info for {name}: {e}") + return storage_info if storage_info else None From 39ec375c73f680558ed047ff2e40a83de3e2b346 Mon Sep 17 00:00:00 2001 From: Danang Massandy Date: Wed, 28 Jan 2026 08:24:15 +0000 Subject: [PATCH 2/3] fix tests --- .../geosight/health/test_readiness.py | 22 ++++++++++--------- 1 file changed, 12 insertions(+), 10 deletions(-) diff --git a/django_project/geosight/health/test_readiness.py b/django_project/geosight/health/test_readiness.py index 73f661d22..b0bfe9a0c 100644 --- a/django_project/geosight/health/test_readiness.py +++ b/django_project/geosight/health/test_readiness.py @@ -14,12 +14,14 @@ __date__ = '22/10/2024' __copyright__ = ('Copyright 2023, Unicef') -from django.test import TestCase, override_settings +from django.test import override_settings from django.urls import reverse from django.core.cache import cache from unittest.mock import patch, MagicMock import tempfile +from core.tests.base_tests import APITestCase + @override_settings( CACHES={ @@ -29,7 +31,7 @@ } } ) -class ReadinessProbeTest(TestCase): +class ReadinessProbeTest(APITestCase): """Test suite for readiness probe endpoint.""" def setUp(self): @@ -64,7 +66,7 @@ def test_database_check_passes(self): self.assertEqual(response.status_code, 200) self.assertTrue(response.json()['checks']['database']) - @patch('health.views.connection') + @patch('geosight.health.views.connection') def test_database_check_fails(self, mock_connection): """Test readiness probe fails when database is down. @@ -90,7 +92,7 @@ def test_redis_check_passes(self): self.assertEqual(response.status_code, 200) self.assertTrue(response.json()['checks']['redis']) - @patch('health.views.cache') + @patch('geosight.health.views.cache') def test_redis_check_fails_on_set(self, mock_cache): """Test readiness probe fails when Redis set operation fails. @@ -103,7 +105,7 @@ def test_redis_check_fails_on_set(self, mock_cache): self.assertEqual(response.json()['status'], 'not ready') self.assertFalse(response.json()['checks']['redis']) - @patch('health.views.cache') + @patch('geosight.health.views.cache') def test_redis_check_fails_on_get(self, mock_cache): """Test readiness probe fails when Redis get returns wrong value. @@ -129,7 +131,7 @@ def test_storage_check_passes_with_low_usage(self): self.assertEqual(response.status_code, 200) self.assertTrue(response.json()['checks']['storage']) - @patch('health.views.shutil.disk_usage') + @patch('geosight.health.views.shutil.disk_usage') def test_storage_check_fails_when_usage_exceeds_threshold( self, mock_disk_usage ): @@ -154,7 +156,7 @@ def test_storage_check_fails_when_usage_exceeds_threshold( self.assertEqual(response.json()['status'], 'not ready') self.assertFalse(response.json()['checks']['storage']) - @patch('health.views.shutil.disk_usage') + @patch('geosight.health.views.shutil.disk_usage') def test_storage_check_passes_at_threshold_boundary(self, mock_disk_usage): """Test storage check passes when exactly at 98% (not exceeding). @@ -171,7 +173,7 @@ def test_storage_check_passes_at_threshold_boundary(self, mock_disk_usage): self.assertEqual(response.status_code, 200) self.assertTrue(response.json()['checks']['storage']) - @patch('health.views.shutil.disk_usage') + @patch('geosight.health.views.shutil.disk_usage') def test_storage_check_with_custom_threshold(self, mock_disk_usage): """Test storage check respects custom threshold setting. @@ -194,8 +196,8 @@ def test_storage_check_with_custom_threshold(self, mock_disk_usage): self.assertEqual(response.status_code, 503) self.assertFalse(response.json()['checks']['storage']) - @patch('health.views.connection') - @patch('health.views.cache') + @patch('geosight.health.views.connection') + @patch('geosight.health.views.cache') def test_multiple_checks_fail(self, mock_cache, mock_connection): """Test readiness probe when multiple checks fail. From 24e1671603cde3ba2992da1d5cc014c2b0d00091 Mon Sep 17 00:00:00 2001 From: Danang Massandy Date: Mon, 2 Feb 2026 08:15:07 +0000 Subject: [PATCH 3/3] add health check command --- .../geosight/health/management/__init__.py | 15 ++++ .../health/management/commands/__init__.py | 15 ++++ .../management/commands/celery_health.py | 83 +++++++++++++++++++ 3 files changed, 113 insertions(+) create mode 100644 django_project/geosight/health/management/__init__.py create mode 100644 django_project/geosight/health/management/commands/__init__.py create mode 100644 django_project/geosight/health/management/commands/celery_health.py diff --git a/django_project/geosight/health/management/__init__.py b/django_project/geosight/health/management/__init__.py new file mode 100644 index 000000000..edbeaf7e3 --- /dev/null +++ b/django_project/geosight/health/management/__init__.py @@ -0,0 +1,15 @@ +# coding=utf-8 +""" +GeoSight is UNICEF's geospatial web-based business intelligence platform. + +Contact : geosight-no-reply@unicef.org + +.. note:: This program is free software; you can redistribute it and/or modify + it under the terms of the GNU Affero General Public License as published by + the Free Software Foundation; either version 3 of the License, or + (at your option) any later version. + +""" +__author__ = 'danang@kartoza.com' +__date__ = '22/10/2024' +__copyright__ = ('Copyright 2023, Unicef') diff --git a/django_project/geosight/health/management/commands/__init__.py b/django_project/geosight/health/management/commands/__init__.py new file mode 100644 index 000000000..edbeaf7e3 --- /dev/null +++ b/django_project/geosight/health/management/commands/__init__.py @@ -0,0 +1,15 @@ +# coding=utf-8 +""" +GeoSight is UNICEF's geospatial web-based business intelligence platform. + +Contact : geosight-no-reply@unicef.org + +.. note:: This program is free software; you can redistribute it and/or modify + it under the terms of the GNU Affero General Public License as published by + the Free Software Foundation; either version 3 of the License, or + (at your option) any later version. + +""" +__author__ = 'danang@kartoza.com' +__date__ = '22/10/2024' +__copyright__ = ('Copyright 2023, Unicef') diff --git a/django_project/geosight/health/management/commands/celery_health.py b/django_project/geosight/health/management/commands/celery_health.py new file mode 100644 index 000000000..dc9f7074b --- /dev/null +++ b/django_project/geosight/health/management/commands/celery_health.py @@ -0,0 +1,83 @@ +# coding=utf-8 +""" +GeoSight is UNICEF's geospatial web-based business intelligence platform. + +Contact : geosight-no-reply@unicef.org + +.. note:: This program is free software; you can redistribute it and/or modify + it under the terms of the GNU Affero General Public License as published by + the Free Software Foundation; either version 3 of the License, or + (at your option) any later version. + +""" +__author__ = 'danang@kartoza.com' +__date__ = '22/10/2024' +__copyright__ = ('Copyright 2023, Unicef') + +from django.core.management.base import BaseCommand +from celery import current_app +import sys + + +class Command(BaseCommand): + """Django management command to check Celery worker health.""" + + help = 'Check Celery worker health and readiness' + + def add_arguments(self, parser): + """Add command line arguments for health check type and timeout.""" + parser.add_argument( + '--check', + type=str, + choices=['liveness', 'readiness'], + required=True, + help='Type of health check to perform' + ) + parser.add_argument( + '--timeout', + type=int, + default=5, + help='Timeout in seconds for the health check' + ) + + def handle(self, *args, **options): + """Handle the health check command.""" + check_type = options['check'] + timeout = options['timeout'] + + try: + if check_type == 'liveness': + self.check_liveness(timeout) + elif check_type == 'readiness': + self.check_readiness(timeout) + + self.stdout.write(self.style.SUCCESS(f'{check_type} check passed')) + sys.exit(0) + except Exception as e: + self.stderr.write( + self.style.ERROR(f'{check_type} check failed: {str(e)}') + ) + sys.exit(1) + + def check_liveness(self, timeout): + """Check if worker is alive and responding.""" + inspect = current_app.control.inspect(timeout=timeout) + pong = inspect.ping() + + if not pong: + raise Exception("No active workers found") + + # pong returns: {'worker@hostname': {'ok': 'pong'}} + if not any('ok' in response for response in pong.values()): + raise Exception("Workers not responding properly") + + def check_readiness(self, timeout): + """Check if worker can connect to broker.""" + conn = current_app.connection() + conn.ensure_connection( + max_retries=3, + interval_start=0, + interval_step=1, + timeout=timeout + ) + conn.release()