From: Jamie Hewland <jhewland@gmail.com>
Date: Tue, 22 Oct 2019 20:01:40 +0000 (+0200)
Subject: asyncio: Wait for the stream to close when closing (#494)
X-Git-Tag: 0.7.6~3
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=88d73de7527e1e4a6d5e7cca1926e201db0b955d;p=thirdparty%2Fhttpx.git

asyncio: Wait for the stream to close when closing (#494)
---

diff --git a/httpx/concurrency/asyncio.py b/httpx/concurrency/asyncio.py
index 70833454..4aeb7ca5 100644
--- a/httpx/concurrency/asyncio.py
+++ b/httpx/concurrency/asyncio.py
@@ -1,6 +1,7 @@
 import asyncio
 import functools
 import ssl
+import sys
 import typing
 from types import TracebackType
 
@@ -170,6 +171,8 @@ class TCPStream(BaseTCPStream):
 
     async def close(self) -> None:
         self.stream_writer.close()
+        if sys.version_info >= (3, 7):
+            await self.stream_writer.wait_closed()
 
 
 class PoolSemaphore(BasePoolSemaphore):
diff --git a/tests/test_concurrency.py b/tests/test_concurrency.py
index 27bbeaf2..8bb933b6 100644
--- a/tests/test_concurrency.py
+++ b/tests/test_concurrency.py
@@ -50,7 +50,21 @@ async def test_start_tls_on_socket_stream(https_server, backend, get_cipher):
         assert get_cipher(stream) is not None
 
         await stream.write(b"GET / HTTP/1.1\r\n\r\n")
-        assert (await stream.read(8192, timeout)).startswith(b"HTTP/1.1 200 OK\r\n")
+
+        # stream.read() only gives us *up to* as much data as we ask for. In order to
+        # cleanly close the stream, we must read until the end of the HTTP response.
+        read = b""
+        ended = False
+        for _ in range(5):  # Try read some (not too large) number of times...
+            read += await stream.read(8192, timeout)
+            # We know we're at the end of the response when we've received the body plus
+            # the terminating CRLFs.
+            if b"Hello, world!" in read and read.endswith(b"\r\n\r\n"):
+                ended = True
+                break
+
+        assert ended
+        assert read.startswith(b"HTTP/1.1 200 OK\r\n")
 
     finally:
         await stream.close()