msd: fixed notifications

1 files changed, 41 insertions, 31 deletions

diff --git a/kvmd/apps/kvmd/msd.py b/kvmd/apps/kvmd/msd.py
index 23671408..019688da 100644
--- a/kvmd/apps/kvmd/msd.py
+++ b/kvmd/apps/kvmd/msd.py
@@ -65,7 +65,7 @@ class MsdOfflineError(MsdOperationError):
         super().__init__("Mass-storage device is not found")
 
 
-class MsdAlreadyConnectedToPcError(MsdOperationError):
+class MsdAlreadyConnectedToServerError(MsdOperationError):
     def __init__(self) -> None:
         super().__init__("Mass-storage is already connected to Server")
 
@@ -278,34 +278,50 @@ class MassStorageDevice:  # pylint: disable=too-many-instance-attributes
     @_msd_working
     @aiotools.atomic
     async def connect_to_kvm(self) -> Dict:
-        with self.__region:
-            if self.__on_kvm:
-                raise MsdAlreadyConnectedToKvmError()
+        notify = False
+        state: Dict = {}
+        try:
+            with self.__region:
+                if self.__on_kvm:
+                    raise MsdAlreadyConnectedToKvmError()
+                notify = True
 
-            gpio.write(self.__target_pin, False)
-            try:
-                await self.__load_device_info()
-            except Exception:
-                if not self.__on_kvm:
-                    gpio.write(self.__target_pin, True)
-                raise
-            self.__on_kvm = True
-            get_logger().info("Mass-storage device switched to KVM: %s", self._device_info)
+                gpio.write(self.__target_pin, False)
+                try:
+                    await self.__load_device_info()
+                except Exception:
+                    if not self.__on_kvm:
+                        gpio.write(self.__target_pin, True)
+                    raise
+                self.__on_kvm = True
+                get_logger().info("Mass-storage device switched to KVM: %s", self._device_info)
 
-        return (await self.__queue_current_state())
+            state = self.get_state()
+            return state
+        finally:
+            if notify:
+                await self.__state_queue.put(state or self.get_state())
 
     @_msd_working
     @aiotools.atomic
     async def connect_to_pc(self) -> Dict:
-        with self.__region:
-            if not self.__on_kvm:
-                raise MsdAlreadyConnectedToPcError()
+        notify = False
+        state: Dict = {}
+        try:
+            with self.__region:
+                if not self.__on_kvm:
+                    raise MsdAlreadyConnectedToServerError()
+                notify = True
 
-            gpio.write(self.__target_pin, True)
-            self.__on_kvm = False
-            get_logger().info("Mass-storage device switched to Server")
+                gpio.write(self.__target_pin, True)
+                self.__on_kvm = False
+                get_logger().info("Mass-storage device switched to Server")
 
-        return (await self.__queue_current_state())
+            state = self.get_state()
+            return state
+        finally:
+            if notify:
+                await self.__state_queue.put(state or self.get_state())
 
     @aiotools.tasked
     @aiotools.atomic
@@ -315,9 +331,7 @@ class MassStorageDevice:  # pylint: disable=too-many-instance-attributes
             with self.__region:
                 if not self._enabled:
                     raise MsdDisabledError()
-
                 notify = True
-                get_logger(0).info("Mass-storage device reset")
 
                 gpio.write(self.__reset_pin, True)
                 await asyncio.sleep(self.__reset_delay)
@@ -327,9 +341,10 @@ class MassStorageDevice:  # pylint: disable=too-many-instance-attributes
                 gpio.write(self.__reset_pin, False)
 
                 await self.__load_device_info()
+                get_logger(0).info("Mass-storage device reset has been successful")
         finally:
             if notify:
-                await self.__queue_current_state()
+                await self.__state_queue.put(self.get_state())
 
     @_msd_working
     @aiotools.atomic
@@ -346,7 +361,7 @@ class MassStorageDevice:  # pylint: disable=too-many-instance-attributes
             self.__region.exit()
             raise
         finally:
-            await self.__queue_current_state()
+            await self.__state_queue.put(self.get_state())
 
     @aiotools.atomic
     async def write_image_info(self, name: str, complete: bool) -> None:
@@ -378,12 +393,7 @@ class MassStorageDevice:  # pylint: disable=too-many-instance-attributes
             await self.__load_device_info()
         finally:
             self.__region.exit()
-            await self.__queue_current_state()
-
-    async def __queue_current_state(self) -> Dict:
-        state = self.get_state()
-        await self.__state_queue.put(state)
-        return state
+            await self.__state_queue.put(self.get_state())
 
     async def __write_to_device_file(self, data: bytes) -> None:
         assert self.__device_file