unified aioregion logic

4 files changed, 15 insertions, 19 deletions

diff --git a/kvmd/kvmd/aioregion.py b/kvmd/kvmd/aioregion.py
index 0e81f82a..c875ad29 100644
--- a/kvmd/kvmd/aioregion.py
+++ b/kvmd/kvmd/aioregion.py
@@ -9,14 +9,15 @@ class RegionIsBusyError(Exception):
 
 
 class AioExclusiveRegion:
-    def __init__(self) -> None:
+    def __init__(self, exc_type: Type[RegionIsBusyError]) -> None:
+        self.__exc_type = exc_type
         self.__busy = False
 
     def enter(self) -> None:
         if not self.__busy:
             self.__busy = True
             return
-        raise RegionIsBusyError()
+        raise self.__exc_type()
 
     def exit(self) -> None:
         self.__busy = False
diff --git a/kvmd/kvmd/atx.py b/kvmd/kvmd/atx.py
index cf155e7e..006c2959 100644
--- a/kvmd/kvmd/atx.py
+++ b/kvmd/kvmd/atx.py
@@ -9,6 +9,11 @@ from . import gpio
 
 
 # =====
+class AtxIsBusy(aioregion.RegionIsBusyError):
+    def __init__(self) -> None:
+        super().__init__("Performing another operation, please try again later")
+
+
 class Atx:
     def __init__(
         self,
@@ -29,7 +34,7 @@ class Atx:
         self.__click_delay = click_delay
         self.__long_click_delay = long_click_delay
 
-        self.__region = aioregion.AioExclusiveRegion()
+        self.__region = aioregion.AioExclusiveRegion(AtxIsBusy)
 
     def get_state(self) -> Dict:
         return {
diff --git a/kvmd/kvmd/msd.py b/kvmd/kvmd/msd.py
index 1dc658c6..8280de40 100644
--- a/kvmd/kvmd/msd.py
+++ b/kvmd/kvmd/msd.py
@@ -49,7 +49,7 @@ class IsNotConnectedToKvmError(MassStorageOperationError):
         super().__init__("Mass-storage is not connected to KVM")
 
 
-class IsBusyError(MassStorageOperationError):
+class IsBusyError(MassStorageOperationError, aioregion.RegionIsBusyError):
     def __init__(self) -> None:
         super().__init__("Mass-storage is busy (write in progress)")
 
@@ -155,8 +155,6 @@ def _explore_device(device_path: str) -> Optional[_MassStorageDeviceInfo]:
 
 def _msd_operated(method: Callable) -> Callable:
     async def wrap(self: "MassStorageDevice", *args: Any, **kwargs: Any) -> Any:
-        if self._device_file:  # pylint: disable=protected-access
-            raise IsBusyError()
         if not self._device_path:  # pylint: disable=protected-access
             IsNotOperationalError()
         return (await method(self, *args, **kwargs))
@@ -179,7 +177,7 @@ class MassStorageDevice:  # pylint: disable=too-many-instance-attributes
         self.__loop = loop
 
         self.__device_info: Optional[_MassStorageDeviceInfo] = None
-        self.__region = aioregion.AioExclusiveRegion()
+        self.__region = aioregion.AioExclusiveRegion(IsBusyError)
         self._device_file: Optional[aiofiles.base.AiofilesContextManager] = None
         self.__written = 0
 
@@ -238,11 +236,11 @@ class MassStorageDevice:  # pylint: disable=too-many-instance-attributes
 
     @_msd_operated
     async def __aenter__(self) -> "MassStorageDevice":
+        self.__region.enter()
         if not self.__device_info:
             raise IsNotConnectedToKvmError()
         self._device_file = await aiofiles.open(self.__device_info.path, mode="w+b", buffering=0)
         self.__written = 0
-        self.__region.enter()
         return self
 
     async def write_image_info(self, name: str, complete: bool) -> None:
diff --git a/kvmd/kvmd/server.py b/kvmd/kvmd/server.py
index 8a32c8c1..d6a690a1 100644
--- a/kvmd/kvmd/server.py
+++ b/kvmd/kvmd/server.py
@@ -58,23 +58,15 @@ class BadRequest(Exception):
     pass
 
 
-class PerformingAnotherOperation(Exception):
-    def __init__(self) -> None:
-        super().__init__("Performing another operation, please try again later")
-
-
 def _wrap_exceptions_for_web(msg: str) -> Callable:
     def make_wrapper(method: Callable) -> Callable:
         async def wrap(self: "Server", request: aiohttp.web.Request) -> aiohttp.web.Response:
             try:
-                try:
-                    return (await method(self, request))
-                except RegionIsBusyError:
-                    raise PerformingAnotherOperation()
+                return (await method(self, request))
+            except RegionIsBusyError as err:
+                return _json_exception(msg, err, 409)
             except (BadRequest, MassStorageOperationError) as err:
                 return _json_exception(msg, err, 400)
-            except PerformingAnotherOperation as err:
-                return _json_exception(msg, err, 409)
         return wrap
     return make_wrapper