Add disk and memory usage checks

pikhovkin · Sep 22, 2021 · 9762527 · 9762527
1 parent 8680d23
commit 9762527
Show file tree

Hide file tree

Showing 5 changed files with 125 additions and 3 deletions.
diff --git a/.github/workflows/python-package.yml b/.github/workflows/python-package.yml
@@ -30,7 +30,7 @@ jobs:
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
-        pip install -e .
+        pip install -e .[psutil]
         pip install "${{ matrix.django }}"
     - name: Test
       run: |

diff --git a/setup.py b/setup.py
@@ -7,7 +7,7 @@
 
 setup(
     name='django-simple-health-check',
-    version='0.2.0',
+    version='0.3.0',
     description='Simple Django health check',
     long_description=open('README.md').read(),
     long_description_content_type='text/markdown',
@@ -21,6 +21,9 @@
     install_requires=[
         'Django>=3.1,<4.0',
     ],
+    extras_require={
+        'psutil': ['psutil'],
+    },
     python_requires='>=3.7.*, <4.0.*',
     license='MIT',
     zip_safe=False,

diff --git a/simple_health_check/checks/ps.py b/simple_health_check/checks/ps.py
@@ -0,0 +1,45 @@
+import locale
+import socket
+
+import psutil
+
+from . import BaseHealthCheck
+from ..exceptions import HealthCheckError
+
+
+host = socket.gethostname()
+
+
+class DiskUsage(BaseHealthCheck):
+    def __init__(self, max_usage_percent: int = 0):
+        super().__init__()
+
+        self.max_usage_percent = max_usage_percent
+
+    def check(self):
+        if not self.max_usage_percent:
+            return
+
+        du = psutil.disk_usage('/')
+        if du.percent >= self.max_usage_percent:
+            raise HealthCheckError(f'{host} {du.percent}% disk usage exceeds {self.max_usage_percent}%')
+
+
+class MemoryUsage(BaseHealthCheck):
+    MB_1 = 1024 * 1024
+
+    def __init__(self, min_memory_mb: int = 0):
+        super().__init__()
+
+        self.min_memory_mb = min_memory_mb
+        self._min_memory_mb = min_memory_mb * self.MB_1
+
+    def check(self):
+        if not self.min_memory_mb:
+            return
+
+        memory = psutil.virtual_memory()
+        if memory.available < self._min_memory_mb:
+            locale.setlocale(locale.LC_ALL, '')
+            available = round(memory.available / self.MB_1, 2)
+            raise HealthCheckError(f'{host} {available:n} MB available RAM below {self.min_memory_mb:d} MB')
diff --git a/simple_health_check/tests.py b/simple_health_check/tests.py
@@ -74,3 +74,75 @@ def test_cache_no_rediness(self):
         response = self.client.get('/readiness/')
         self.assertTrue(response.status_code == 500)
         self.assertTrue(response.content == b'down')
+
+    @override_settings(
+        SIMPLE_HEALTH_CHECKS={
+            'simple_health_check.checks.ps.DiskUsage': None,
+        },
+    )
+    def test_ps_disk_usage_no_value(self):
+        apps.get_app_config('simple_health_check').register_checks()
+
+        response = self.client.get('/readiness/')
+        self.assertTrue(response.status_code == 200)
+        self.assertTrue(response.content == b'ok')
+
+    @override_settings(
+        SIMPLE_HEALTH_CHECKS={
+            'simple_health_check.checks.ps.DiskUsage': dict(max_usage_percent=99),
+        },
+    )
+    def test_ps_disk_usage(self):
+        apps.get_app_config('simple_health_check').register_checks()
+
+        response = self.client.get('/readiness/')
+        self.assertTrue(response.status_code == 200)
+        self.assertTrue(response.content == b'ok')
+
+    @override_settings(
+        SIMPLE_HEALTH_CHECKS={
+            'simple_health_check.checks.ps.DiskUsage': dict(max_usage_percent=0.001),
+        },
+    )
+    def test_ps_disk_usage_no_rediness(self):
+        apps.get_app_config('simple_health_check').register_checks()
+
+        response = self.client.get('/readiness/')
+        self.assertTrue(response.status_code == 500)
+        self.assertTrue(response.content == b'down')
+
+    @override_settings(
+        SIMPLE_HEALTH_CHECKS={
+            'simple_health_check.checks.ps.MemoryUsage': None,
+        },
+    )
+    def test_ps_memory_usage_no_value(self):
+        apps.get_app_config('simple_health_check').register_checks()
+
+        response = self.client.get('/readiness/')
+        self.assertTrue(response.status_code == 200)
+        self.assertTrue(response.content == b'ok')
+
+    @override_settings(
+        SIMPLE_HEALTH_CHECKS={
+            'simple_health_check.checks.ps.MemoryUsage': dict(min_memory_mb=10),
+        },
+    )
+    def test_ps_memory_usage(self):
+        apps.get_app_config('simple_health_check').register_checks()
+
+        response = self.client.get('/readiness/')
+        self.assertTrue(response.status_code == 200)
+        self.assertTrue(response.content == b'ok')
+
+    @override_settings(
+        SIMPLE_HEALTH_CHECKS={
+            'simple_health_check.checks.ps.MemoryUsage': dict(min_memory_mb=100_000),
+        },
+    )
+    def test_ps_memory_usage_no_rediness(self):
+        apps.get_app_config('simple_health_check').register_checks()
+
+        response = self.client.get('/readiness/')
+        self.assertTrue(response.status_code == 500)
+        self.assertTrue(response.content == b'down')
diff --git a/tox.ini b/tox.ini
@@ -4,7 +4,8 @@ skip_missing_interpreters = True
 envlist = py{37,38,39,310}-django{31,32}
 
 [testenv]
-commands = python manage.py test
+extras =
+    psutil
 deps =
     django31: Django>=3.1,<3.2
     django32: Django>=3.2,<3.3
@@ -14,3 +15,4 @@ basepython =
     py38: {env:TOX_PYTHON_38}
     py39: {env:TOX_PYTHON_39}
     py310: {env:TOX_PYTHON_310}
+commands = python manage.py test