From 441460d81ff2b1faee7d044d859896f754361b93 Mon Sep 17 00:00:00 2001
From: Damien George <damien@micropython.org>
Date: Sat, 25 Jul 2020 23:05:41 +1000
Subject: [PATCH] extmod/uasyncio: Add StreamReader.readexactly(n) method.

It raises on EOFError instead of an IncompleteReadError (which is what
CPython does).  But the latter is derived from EOFError so code compatible
with MicroPython and CPython can be written by catching EOFError (eg see
included test).

Fixes issue #6156.

Signed-off-by: Damien George <damien@micropython.org>
---
 extmod/uasyncio/stream.py                     | 12 ++++
 tests/multi_net/uasyncio_tcp_readexactly.py   | 68 +++++++++++++++++++
 .../multi_net/uasyncio_tcp_readexactly.py.exp | 10 +++
 3 files changed, 90 insertions(+)
 create mode 100644 tests/multi_net/uasyncio_tcp_readexactly.py
 create mode 100644 tests/multi_net/uasyncio_tcp_readexactly.py.exp

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''