diff --git a/extmod/uasyncio/stream.py b/extmod/uasyncio/stream.py index 2a1efd1a17..b6d787e4f0 100644 --- a/extmod/uasyncio/stream.py +++ b/extmod/uasyncio/stream.py @@ -30,6 +30,18 @@ class Stream: yield core._io_queue.queue_read(self.s) return self.s.read(n) + async def readexactly(self, n): + r = b"" + while n: + yield core._io_queue.queue_read(self.s) + r2 = self.s.read(n) + if r2 is not None: + if not len(r2): + raise EOFError + r += r2 + n -= len(r2) + return r + async def readline(self): l = b"" while True: diff --git a/tests/multi_net/uasyncio_tcp_readexactly.py b/tests/multi_net/uasyncio_tcp_readexactly.py new file mode 100644 index 0000000000..71d8c6d0e9 --- /dev/null +++ b/tests/multi_net/uasyncio_tcp_readexactly.py @@ -0,0 +1,68 @@ +# Test uasyncio stream readexactly() method using TCP server/client + +try: + import uasyncio as asyncio +except ImportError: + try: + import asyncio + except ImportError: + print("SKIP") + raise SystemExit + +PORT = 8000 + + +async def handle_connection(reader, writer): + writer.write(b"a") + await writer.drain() + + # Split the first 2 bytes up so the client must wait for the second one + await asyncio.sleep(0.1) + + writer.write(b"b") + await writer.drain() + + writer.write(b"c") + await writer.drain() + + writer.write(b"d") + await writer.drain() + + print("close") + writer.close() + await writer.wait_closed() + + print("done") + ev.set() + + +async def tcp_server(): + global ev + ev = asyncio.Event() + server = await asyncio.start_server(handle_connection, "0.0.0.0", PORT) + print("server running") + multitest.next() + async with server: + await asyncio.wait_for(ev.wait(), 10) + + +async def tcp_client(): + reader, writer = await asyncio.open_connection(IP, PORT) + print(await reader.readexactly(2)) + print(await reader.readexactly(0)) + print(await reader.readexactly(1)) + try: + print(await reader.readexactly(2)) + except EOFError as er: + print("EOFError") + print(await reader.readexactly(0)) + + +def instance0(): + multitest.globals(IP=multitest.get_network_ip()) + asyncio.run(tcp_server()) + + +def instance1(): + multitest.next() + asyncio.run(tcp_client()) diff --git a/tests/multi_net/uasyncio_tcp_readexactly.py.exp b/tests/multi_net/uasyncio_tcp_readexactly.py.exp new file mode 100644 index 0000000000..65ce6d628d --- /dev/null +++ b/tests/multi_net/uasyncio_tcp_readexactly.py.exp @@ -0,0 +1,10 @@ +--- instance0 --- +server running +close +done +--- instance1 --- +b'ab' +b'' +b'c' +EOFError +b''