Clean up failed storage probe files

Author: M-Hassan-Raza

health_check/checks.py

@@ -253,16 +253,23 @@ def get_file_name(self):
     def get_file_content(self):
         return f"# generated by health_check.Storage at {datetime.datetime.now().timestamp()}".encode()
 
-    def check_save(self, file_name, file_content):
-        # save the file
-        file_name = self.storage.save(file_name, ContentFile(content=file_content))
+    def check_written_file(self, file_name, file_content):
         # read the file and compare
         if not self.storage.exists(file_name):
             raise ServiceUnavailable("File does not exist")
         with self.storage.open(file_name) as f:
             if not f.read() == file_content:
                 raise ServiceUnavailable("File content does not match")
-        return file_name
+
+    def cleanup_file(self, file_name):
+        try:
+            self.storage.delete(file_name)
+        except Exception:
+            logger.warning(
+                "Failed to clean up storage probe file %r after a failed health check.",
+                file_name,
+                exc_info=True,
+            )
 
     def check_delete(self, file_name):
         # delete the file and make sure it is gone
@@ -274,5 +281,10 @@ def run(self):
         # write the file to the storage backend
         file_name = self.get_file_name()
         file_content = self.get_file_content()
-        file_name = self.check_save(file_name, file_content)
+        file_name = self.storage.save(file_name, ContentFile(content=file_content))
+        try:
+            self.check_written_file(file_name, file_content)
+        except Exception:
+            self.cleanup_file(file_name)
+            raise
         self.check_delete(file_name)

tests/test_checks.py

@@ -483,37 +483,55 @@ async def test_check_status__not_deleted(self):
             assert "File was not deleted" in str(result.error)
 
     @pytest.mark.asyncio
-    async def test_check_status__file_not_saved(self):
-        """Raise ServiceUnavailable when file does not exist after save."""
+    @pytest.mark.parametrize(
+        ("exists_result", "read_result", "expected_message"),
+        [
+            (False, None, "does not exist"),
+            (True, b"wrong content", "does not match"),
+        ],
+    )
+    async def test_check_status__failed_validation_cleans_up_saved_file(
+        self, exists_result, read_result, expected_message
+    ):
+        """Delete the saved probe file when validation fails after write."""
         with mock.patch("health_check.checks.storages") as mock_storages:
             mock_storage = mock.MagicMock()
             mock_storages.__getitem__.return_value = mock_storage
             mock_storage.save.return_value = "test-file.txt"
-            mock_storage.exists.return_value = False
+            mock_storage.exists.return_value = exists_result
+            if read_result is not None:
+                mock_file = mock.MagicMock()
+                mock_file.read.return_value = read_result
+                mock_storage.open.return_value.__enter__.return_value = mock_file
 
             check = Storage()
             result = await check.get_result()
             assert result.error is not None
             assert isinstance(result.error, ServiceUnavailable)
-            assert "does not exist" in str(result.error)
+            assert expected_message in str(result.error)
+            mock_storage.delete.assert_called_once_with("test-file.txt")
 
     @pytest.mark.asyncio
-    async def test_check_status__file_content_mismatch(self):
-        """Raise ServiceUnavailable when file content does not match."""
-        with mock.patch("health_check.checks.storages") as mock_storages:
+    async def test_check_status__cleanup_failure_does_not_mask_validation_error(self):
+        """Log cleanup failure and return the original validation error."""
+        with (
+            mock.patch("health_check.checks.storages") as mock_storages,
+            mock.patch("health_check.checks.logger") as mock_logger,
+        ):
             mock_storage = mock.MagicMock()
             mock_storages.__getitem__.return_value = mock_storage
             mock_storage.save.return_value = "test-file.txt"
-            mock_storage.exists.return_value = True
-            mock_file = mock.MagicMock()
-            mock_file.read.return_value = b"wrong content"
-            mock_storage.open.return_value.__enter__.return_value = mock_file
+            mock_storage.exists.return_value = False
+            mock_storage.delete.side_effect = OSError("delete failed")
 
             check = Storage()
             result = await check.get_result()
+
             assert result.error is not None
             assert isinstance(result.error, ServiceUnavailable)
-            assert "does not match" in str(result.error)
+            assert "does not exist" in str(result.error)
+            mock_storage.delete.assert_called_once_with("test-file.txt")
+            mock_logger.warning.assert_called_once()
 
     @pytest.mark.asyncio
     async def test_check_status__service_unavailable_passthrough(self):