diff options
Diffstat (limited to 'kvmd/apps/vnc')
-rw-r--r-- | kvmd/apps/vnc/kvmd.py | 3 | ||||
-rw-r--r-- | kvmd/apps/vnc/streamer.py | 4 |
2 files changed, 5 insertions, 2 deletions
diff --git a/kvmd/apps/vnc/kvmd.py b/kvmd/apps/vnc/kvmd.py index 5d130b02..6d492d5e 100644 --- a/kvmd/apps/vnc/kvmd.py +++ b/kvmd/apps/vnc/kvmd.py @@ -23,6 +23,7 @@ import contextlib from typing import Dict +from typing import AsyncGenerator import aiohttp @@ -64,7 +65,7 @@ class KvmdClient: raise @contextlib.asynccontextmanager - async def ws(self, user: str, passwd: str) -> aiohttp.ClientWebSocketResponse: # pylint: disable=invalid-name + async def ws(self, user: str, passwd: str) -> AsyncGenerator[aiohttp.ClientWebSocketResponse, None]: async with self.__make_session(user, passwd) as session: async with session.ws_connect( url=f"http://{self.__host}:{self.__port}/ws", diff --git a/kvmd/apps/vnc/streamer.py b/kvmd/apps/vnc/streamer.py index 62fce849..fd2232ee 100644 --- a/kvmd/apps/vnc/streamer.py +++ b/kvmd/apps/vnc/streamer.py @@ -61,7 +61,9 @@ class StreamerClient: response.raise_for_status() reader = aiohttp.MultipartReader.from_response(response) while True: - frame = await reader.next() + frame = await reader.next() # pylint: disable=not-callable + if not isinstance(frame, aiohttp.BodyPartReader): + raise RuntimeError("Expected body part") yield ( (frame.headers["X-UStreamer-Online"] == "true"), int(frame.headers["X-UStreamer-Width"]), |