/msd/write_remote handle

1 files changed, 2 insertions, 5 deletions

diff --git a/kvmd/plugins/msd/relay/__init__.py b/kvmd/plugins/msd/relay/__init__.py
index 82de6596..a608e535 100644
--- a/kvmd/plugins/msd/relay/__init__.py
+++ b/kvmd/plugins/msd/relay/__init__.py
@@ -208,7 +208,7 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
                 self.__connected = connected
 
     @contextlib.asynccontextmanager
-    async def write_image(self, name: str, size: int) -> AsyncGenerator[None, None]:
+    async def write_image(self, name: str, size: int) -> AsyncGenerator[int, None]:
         async with self.__working():
             async with self.__region:
                 try:
@@ -220,15 +220,12 @@ class Plugin(BaseMsd):  # pylint: disable=too-many-instance-attributes
 
                     await self.__write_image_info(False)
                     await self.__notifier.notify()
-                    yield
+                    yield self.__upload_chunk_size
                     await self.__write_image_info(True)
                 finally:
                     await self.__close_device_writer()
                     await self.__load_device_info()
 
-    def get_upload_chunk_size(self) -> int:
-        return self.__upload_chunk_size
-
     async def write_image_chunk(self, chunk: bytes) -> int:
         assert self.__device_writer
         return (await self.__device_writer.write(chunk))