diff --git a/pytec/pytec/aioclient.py b/pytec/pytec/aioclient.py index ab4fb6a..aefa07e 100644 --- a/pytec/pytec/aioclient.py +++ b/pytec/pytec/aioclient.py @@ -74,12 +74,16 @@ class Client: chunk = await self._reader.readline() return chunk.decode('utf-8', errors='ignore') + async def _read_write(self, command): + self._writer.write(((" ".join(command)).strip() + "\n").encode('utf-8')) + await self._writer.drain() + + return await self._read_line() + async def _command(self, *command): async with self._command_lock: - self._writer.write(((" ".join(command)).strip() + "\n").encode('utf-8')) - await self._writer.drain() - - line = await self._read_line() + # protect the read-write process from being cancelled midway + line = await asyncio.shield(self._read_write(command)) response = json.loads(line) logging.debug(f"{command}: {response}")