refactoring

4 files changed, 37 insertions, 53 deletions

diff --git a/kvmd/plugins/msd/__init__.py b/kvmd/plugins/msd/__init__.py
index 6edc6ecb..bd397657 100644
--- a/kvmd/plugins/msd/__init__.py
+++ b/kvmd/plugins/msd/__init__.py
@@ -109,10 +109,7 @@ class BaseMsd(BasePlugin):
     async def set_params(self, name: Optional[str]=None, cdrom: Optional[bool]=None) -> None:
         raise NotImplementedError()
 
-    async def connect(self) -> None:
-        raise NotImplementedError()
-
-    async def disconnect(self) -> None:
+    async def set_connected(self, connected: bool) -> None:
         raise NotImplementedError()
 
     @contextlib.asynccontextmanager
diff --git a/kvmd/plugins/msd/disabled.py b/kvmd/plugins/msd/disabled.py
index 179af7da..83746b4c 100644
--- a/kvmd/plugins/msd/disabled.py
+++ b/kvmd/plugins/msd/disabled.py
@@ -66,10 +66,7 @@ class Plugin(BaseMsd):
     async def set_params(self, name: Optional[str]=None, cdrom: Optional[bool]=None) -> None:
         raise MsdDisabledError()
 
-    async def connect(self) -> None:
-        raise MsdDisabledError()
-
-    async def disconnect(self) -> None:
+    async def set_connected(self, connected: bool) -> None:
         raise MsdDisabledError()
 
     @contextlib.asynccontextmanager
diff --git a/kvmd/plugins/msd/otg/__init__.py b/kvmd/plugins/msd/otg/__init__.py
index a4fd3919..2b31ea53 100644
--- a/kvmd/plugins/msd/otg/__init__.py
+++ b/kvmd/plugins/msd/otg/__init__.py
@@ -255,36 +255,32 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
                 self.__state.vd.cdrom = cdrom
 
     @aiotools.atomic
-    async def connect(self) -> None:
+    async def set_connected(self, connected: bool) -> None:
         async with self.__state.busy():
             assert self.__state.vd
+            if connected:
+                if self.__state.vd.connected or self.__drive.get_image_path():
+                    raise MsdConnectedError()
+                if self.__state.vd.image is None:
+                    raise MsdImageNotSelected()
 
-            if self.__state.vd.connected or self.__drive.get_image_path():
-                raise MsdConnectedError()
-            if self.__state.vd.image is None:
-                raise MsdImageNotSelected()
-
-            assert self.__state.vd.image.in_storage
+                assert self.__state.vd.image.in_storage
 
-            if not os.path.exists(self.__state.vd.image.path):
-                raise MsdUnknownImageError()
+                if not os.path.exists(self.__state.vd.image.path):
+                    raise MsdUnknownImageError()
 
-            await self.__unlock_drive()
-            self.__drive.set_cdrom_flag(self.__state.vd.cdrom)
-            self.__drive.set_image_path(self.__state.vd.image.path)
-            self.__state.vd.connected = True
+                await self.__unlock_drive()
+                self.__drive.set_cdrom_flag(self.__state.vd.cdrom)
+                self.__drive.set_image_path(self.__state.vd.image.path)
 
-    @aiotools.atomic
-    async def disconnect(self) -> None:
-        async with self.__state.busy():
-            assert self.__state.vd
+            else:
+                if not (self.__state.vd.connected or self.__drive.get_image_path()):
+                    raise MsdDisconnectedError()
 
-            if not (self.__state.vd.connected or self.__drive.get_image_path()):
-                raise MsdDisconnectedError()
+                await self.__unlock_drive()
+                self.__drive.set_image_path("")
 
-            await self.__unlock_drive()
-            self.__drive.set_image_path("")
-            self.__state.vd.connected = False
+            self.__state.vd.connected = connected
 
     @contextlib.asynccontextmanager
     async def write_image(self, name: str) -> AsyncGenerator[None, None]:
diff --git a/kvmd/plugins/msd/relay/__init__.py b/kvmd/plugins/msd/relay/__init__.py
index ea87168c..bbb2b647 100644
--- a/kvmd/plugins/msd/relay/__init__.py
+++ b/kvmd/plugins/msd/relay/__init__.py
@@ -176,32 +176,26 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
                 raise MsdCdromNotSupported()
 
     @aiotools.atomic
-    async def connect(self) -> None:
+    async def set_connected(self, connected: bool) -> None:
         async with self.__working():
             async with self.__region:
-                if self.__connected:
-                    raise MsdConnectedError()
-
-                self.__gpio.switch_to_server()
-                self.__connected = True
-                get_logger(0).info("MSD switched to Server")
-
-    @aiotools.atomic
-    async def disconnect(self) -> None:
-        async with self.__working():
-            async with self.__region:
-                if not self.__connected:
-                    raise MsdDisconnectedError()
-
-                self.__gpio.switch_to_local()
-                try:
-                    await self.__load_device_info()
-                except Exception:
+                if connected:
                     if self.__connected:
-                        self.__gpio.switch_to_server()
-                    raise
-                self.__connected = False
-                get_logger(0).info("MSD switched to KVM: %s", self.__device_info)
+                        raise MsdConnectedError()
+                    self.__gpio.switch_to_server()
+                    get_logger(0).info("MSD switched to Server")
+                else:
+                    if not self.__connected:
+                        raise MsdDisconnectedError()
+                    self.__gpio.switch_to_local()
+                    try:
+                        await self.__load_device_info()
+                    except Exception:
+                        if self.__connected:
+                            self.__gpio.switch_to_server()
+                        raise
+                    get_logger(0).info("MSD switched to KVM: %s", self.__device_info)
+                self.__connected = connected
 
     @contextlib.asynccontextmanager
     async def write_image(self, name: str) -> AsyncGenerator[None, None]: