From 43c158b1bb5c2b61d2c6c8af3ba3ad7d7d1d7beb Mon Sep 17 00:00:00 2001
From: Christian Hartung <hartung@live.com>
Date: Fri, 11 Oct 2024 17:48:50 -0300
Subject: [PATCH] feat: add StorageHealthCheck

---
 django_healthy/health_checks/storage.py | 38 +++++++++++++++
 tests/health_checks/test_storage.py     | 64 +++++++++++++++++++++++++
 2 files changed, 102 insertions(+)
 create mode 100644 django_healthy/health_checks/storage.py
 create mode 100644 tests/health_checks/test_storage.py

diff --git a/django_healthy/health_checks/storage.py b/django_healthy/health_checks/storage.py
new file mode 100644
index 0000000..f224f97
--- /dev/null
+++ b/django_healthy/health_checks/storage.py
@@ -0,0 +1,38 @@
+from __future__ import annotations
+
+from io import StringIO
+
+from asgiref.sync import sync_to_async
+from django.core.files.storage import Storage, storages
+from django.utils.crypto import get_random_string
+
+from .base import HealthCheck, HealthCheckResult
+
+
+class StorageHealthCheck(HealthCheck):
+    __slots__: tuple[str, ...] = ("alias", "filename")
+
+    def __init__(self, alias: str = "default", filename: str = "healthy_test_file.txt"):
+        self.alias = alias
+        self.filename = filename
+
+    async def check_health(self) -> HealthCheckResult:
+        storage: Storage = storages[self.alias]
+        given: str = get_random_string(100)
+
+        try:
+            filename = await sync_to_async(storage.get_available_name)(self.filename)
+            await sync_to_async(storage.save)(filename, StringIO(given))
+
+            exists = await sync_to_async(storage.exists)(filename)
+            if not exists:
+                return HealthCheckResult.degraded(description="File missing", data={"filename": filename})
+
+            await sync_to_async(storage.delete)(filename)
+            exists = await sync_to_async(storage.exists)(filename)
+            if exists:
+                return HealthCheckResult.degraded(description="Could not delete file", data={"filename": filename})
+        except Exception as exc:  # noqa: BLE001
+            return HealthCheckResult.unhealthy(exception=exc)
+        else:
+            return HealthCheckResult.healthy()
diff --git a/tests/health_checks/test_storage.py b/tests/health_checks/test_storage.py
new file mode 100644
index 0000000..0e3068d
--- /dev/null
+++ b/tests/health_checks/test_storage.py
@@ -0,0 +1,64 @@
+from unittest import mock
+
+import pytest
+from django.core.files.storage import storages
+
+from django_healthy.health_checks import HealthStatus
+from django_healthy.health_checks.storage import StorageHealthCheck
+
+
+@pytest.mark.asyncio
+class TestStorageHealthCheck:
+    async def test_with_working_storage(self):
+        health_check = StorageHealthCheck()
+
+        got = await health_check.check_health()
+
+        assert got.status == HealthStatus.HEALTHY
+
+    async def test_without_saving(self):
+        health_check = StorageHealthCheck()
+        storage = storages[health_check.alias]
+
+        with mock.patch.object(storage, "save"):
+            got = await health_check.check_health()
+
+        assert got.status == HealthStatus.DEGRADED
+        assert "filename" in got.data
+
+    async def test_without_deleting(self):
+        health_check = StorageHealthCheck()
+        storage = storages[health_check.alias]
+
+        with mock.patch.object(storage, "delete"):
+            got = await health_check.check_health()
+
+        assert got.status == HealthStatus.DEGRADED
+        assert "filename" in got.data
+
+    async def test_with_save_error(self):
+        health_check = StorageHealthCheck()
+        storage = storages[health_check.alias]
+
+        with mock.patch.object(storage, "save", side_effect=Exception):
+            got = await health_check.check_health()
+
+        assert got.status == HealthStatus.UNHEALTHY
+
+    async def test_with_delete_error(self):
+        health_check = StorageHealthCheck()
+        storage = storages[health_check.alias]
+
+        with mock.patch.object(storage, "delete", side_effect=Exception):
+            got = await health_check.check_health()
+
+        assert got.status == HealthStatus.UNHEALTHY
+
+    async def test_with_exists_error(self):
+        health_check = StorageHealthCheck()
+        storage = storages[health_check.alias]
+
+        with mock.patch.object(storage, "exists", side_effect=Exception):
+            got = await health_check.check_health()
+
+        assert got.status == HealthStatus.UNHEALTHY