From 3b48d3be0d8ceddfd75ee6bd27e5a4cf41fdd21d Mon Sep 17 00:00:00 2001
From: germanhydrogen <admin@germanhydrogen.org>
Date: Wed, 27 Dec 2023 09:46:53 +0100
Subject: [PATCH] Added exception handler image processor decorator

---
 pyobs/images/processors/__init__.py           |  4 ++
 pyobs/images/processors/exceptionhandler.py   | 41 ++++++++++++++++
 .../processors/test_exceptionhandler.py       | 49 +++++++++++++++++++
 3 files changed, 94 insertions(+)
 create mode 100644 pyobs/images/processors/exceptionhandler.py
 create mode 100644 tests/images/processors/test_exceptionhandler.py

diff --git a/pyobs/images/processors/__init__.py b/pyobs/images/processors/__init__.py
index 018760369..e0c041dd7 100644
--- a/pyobs/images/processors/__init__.py
+++ b/pyobs/images/processors/__init__.py
@@ -2,3 +2,7 @@
 Image Processors (pyobs.images.processors)
 ------------------------------------------
 """
+
+from .exceptionhandler import (ExceptionHandler)
+
+__all__ = ["ExceptionHandler"]
diff --git a/pyobs/images/processors/exceptionhandler.py b/pyobs/images/processors/exceptionhandler.py
new file mode 100644
index 000000000..77a4ecbb1
--- /dev/null
+++ b/pyobs/images/processors/exceptionhandler.py
@@ -0,0 +1,41 @@
+import logging
+from typing import Union, Dict, Any, Optional
+
+from pyobs.images import ImageProcessor, Image
+from pyobs.object import get_object
+
+import pyobs.utils.exceptions as exc
+log = logging.getLogger(__name__)
+
+
+class ExceptionHandler(ImageProcessor):
+
+    __module__ = "pyobs.images.processors"
+
+    def __init__(self, processor: Union[ImageProcessor, Dict[str, Any]], error_header: Optional[str] = None) -> None:
+        super().__init__()
+
+        self._processor: ImageProcessor = get_object(processor, ImageProcessor)  # type: ignore
+        self._error_header = error_header
+
+    async def __call__(self, image: Image) -> Image:
+        try:
+            return await self._processor(image)
+        except exc.ImageError as e:
+            return await self._handle_error(image, e)
+
+    async def _handle_error(self, image: Image, error: exc.ImageError) -> Image:
+        log.warning(error.message)
+
+        output_image = image.copy()
+
+        if self._error_header is not None:
+            output_image.header[self._error_header] = 1
+
+        return output_image
+
+    async def reset(self) -> None:
+        await self._processor.reset()
+
+
+__all__ = ["ExceptionHandler"]
diff --git a/tests/images/processors/test_exceptionhandler.py b/tests/images/processors/test_exceptionhandler.py
new file mode 100644
index 000000000..b63ff978c
--- /dev/null
+++ b/tests/images/processors/test_exceptionhandler.py
@@ -0,0 +1,49 @@
+import logging
+from unittest.mock import Mock
+
+import pytest
+
+from pyobs.images import ImageProcessor, Image
+from pyobs.images.processors import ExceptionHandler
+
+import pyobs.utils.exceptions as exc
+
+class MockImageProcessor(ImageProcessor):
+
+    async def __call__(self, image: Image) -> Image:
+        return image
+
+
+@pytest.mark.asyncio
+async def test_exception_handler_no_exception() -> None:
+    image = Image()
+    exception_handler = ExceptionHandler(MockImageProcessor())
+
+    result = await exception_handler(image)
+
+    assert image == result
+
+
+@pytest.mark.asyncio
+async def test_exception_handler_exception(caplog) -> None:
+    image = Image()
+    exception_handler = ExceptionHandler(MockImageProcessor())
+    exception_handler._processor = Mock(side_effect=exc.ImageError("Some error"))
+
+    with caplog.at_level(logging.WARNING):
+        result = await exception_handler(image)
+
+    assert caplog.messages[0] == "Some error"
+
+    assert image is not result
+
+
+@pytest.mark.asyncio
+async def test_exception_handler_exception_w_header() -> None:
+    image = Image()
+    exception_handler = ExceptionHandler(MockImageProcessor(), "WCSERR")
+    exception_handler._processor = Mock(side_effect=exc.ImageError("Some error"))
+
+    result = await exception_handler(image)
+
+    assert result.header["WCSERR"] == 1
\ No newline at end of file