sanic/tests/client.py
2022-06-27 11:19:26 +03:00

48 lines
1.2 KiB
Python

import asyncio
from textwrap import dedent
from typing import AnyStr
class RawClient:
CRLF = b"\r\n"
def __init__(self, host: str, port: int):
self.reader = None
self.writer = None
self.host = host
self.port = port
async def connect(self):
self.reader, self.writer = await asyncio.open_connection(
self.host, self.port
)
async def close(self):
self.writer.close()
await self.writer.wait_closed()
async def send(self, message: AnyStr):
if isinstance(message, str):
msg = self._clean(message).encode("utf-8")
else:
msg = message
await self._send(msg)
async def _send(self, message: bytes):
if not self.writer:
raise Exception("No open write stream")
self.writer.write(message)
async def recv(self, nbytes: int = -1) -> bytes:
if not self.reader:
raise Exception("No open read stream")
return await self.reader.read(nbytes)
def _clean(self, message: str) -> str:
return (
dedent(message)
.lstrip("\n")
.replace("\n", self.CRLF.decode("utf-8"))
)