refactored msd writer api

1 files changed, 6 insertions, 4 deletions

diff --git a/kvmd/apps/kvmd/api/msd.py b/kvmd/apps/kvmd/api/msd.py
index b6f21aad..259a0a53 100644
--- a/kvmd/apps/kvmd/api/msd.py
+++ b/kvmd/apps/kvmd/api/msd.py
@@ -100,12 +100,13 @@ class MsdApi:
         size = valid_int_f0(request.content_length)
         remove_incomplete = self.__get_remove_incomplete(request)
         written = 0
-        async with self.__msd.write_image(name, size, remove_incomplete) as chunk_size:
+        async with self.__msd.write_image(name, size, remove_incomplete) as writer:
+            chunk_size = writer.get_chunk_size()
             while True:
                 chunk = await request.content.read(chunk_size)
                 if not chunk:
                     break
-                written = await self.__msd.write_image_chunk(chunk)
+                written = await writer.write_chunk(chunk)
         return make_json_response(self.__make_write_info(name, size, written))
 
     @exposed_http("POST", "/msd/write_remote")
@@ -139,12 +140,13 @@ class MsdApi:
                 size = valid_int_f0(remote.content_length)
 
                 get_logger(0).info("Downloading image %r as %r to MSD ...", url, name)
-                async with self.__msd.write_image(name, size, remove_incomplete) as chunk_size:
+                async with self.__msd.write_image(name, size, remove_incomplete) as writer:
+                    chunk_size = writer.get_chunk_size()
                     response = await start_streaming(request, "application/x-ndjson")
                     await stream_write_info()
                     last_report_ts = 0
                     async for chunk in remote.content.iter_chunked(chunk_size):
-                        written = await self.__msd.write_image_chunk(chunk)
+                        written = await writer.write_chunk(chunk)
                         now = int(time.time())
                         if last_report_ts + 1 < now:
                             await stream_write_info()