From c652b9224ecdc513715438c9b81d9d8ce501ec89 Mon Sep 17 00:00:00 2001 From: Dobrovolsky Bogdan Date: Fri, 6 Sep 2019 15:45:52 +0300 Subject: [PATCH] Close all outstanding responses on client close --- httpx/dispatch/connection_pool.py | 18 ++++++----- tests/client/test_async_client.py | 20 ++++++++++--- tests/client/test_client.py | 40 ++++++++++++++++--------- tests/dispatch/test_connection_pools.py | 11 +++++++ 4 files changed, 64 insertions(+), 25 deletions(-) diff --git a/httpx/dispatch/connection_pool.py b/httpx/dispatch/connection_pool.py index befa88cb40..a553915df0 100644 --- a/httpx/dispatch/connection_pool.py +++ b/httpx/dispatch/connection_pool.py @@ -1,4 +1,5 @@ import typing +from itertools import chain from ..concurrency.asyncio import AsyncioBackend from ..concurrency.base import ConcurrencyBackend @@ -62,10 +63,11 @@ def add(self, connection: HTTPConnection) -> None: self.by_origin[connection.origin] = {connection: 0.0} def remove(self, connection: HTTPConnection) -> None: - del self.all[connection] - del self.by_origin[connection.origin][connection] - if not self.by_origin[connection.origin]: - del self.by_origin[connection.origin] + if connection in self.all: + del self.all[connection] + del self.by_origin[connection.origin][connection] + if not self.by_origin[connection.origin]: + del self.by_origin[connection.origin] def clear(self) -> None: self.all.clear() @@ -175,7 +177,9 @@ async def release_connection(self, connection: HTTPConnection) -> None: async def close(self) -> None: self.is_closed = True - connections = list(self.keepalive_connections) - self.keepalive_connections.clear() - for connection in connections: + + for connection in chain(self.keepalive_connections, self.active_connections): await connection.close() + + self.keepalive_connections.clear() + self.active_connections.clear() diff --git a/tests/client/test_async_client.py b/tests/client/test_async_client.py index c028c1b3fb..a7a78ed362 100644 --- a/tests/client/test_async_client.py +++ b/tests/client/test_async_client.py @@ -1,4 +1,5 @@ import pytest +from h11 import RemoteProtocolError import httpx @@ -43,13 +44,24 @@ async def test_post_json(server, backend): assert response.status_code == 200 +async def test_stream_manger_release_resources(server, backend): + async with httpx.AsyncClient(backend=backend) as http: + response = await http.post( + server.url.copy_with(path="/echo_body"), stream=True, data=b"*" * 409600 + ) + + assert response.status_code == 200 + with pytest.raises(RemoteProtocolError): + await response.read() + + async def test_stream_response(server, backend): async with httpx.AsyncClient(backend=backend) as client: response = await client.request("GET", server.url, stream=True) - assert response.status_code == 200 - body = await response.read() - assert body == b"Hello, world!" - assert response.content == b"Hello, world!" + assert response.status_code == 200 + body = await response.read() + assert body == b"Hello, world!" + assert response.content == b"Hello, world!" async def test_access_content_stream_response(server, backend): diff --git a/tests/client/test_client.py b/tests/client/test_client.py index 73979209af..3c79fe7660 100644 --- a/tests/client/test_client.py +++ b/tests/client/test_client.py @@ -1,4 +1,5 @@ import pytest +from h11 import RemoteProtocolError import httpx @@ -33,33 +34,44 @@ def test_post_json(server): assert response.reason_phrase == "OK" +def test_stream_manger_release_resources(server): + with httpx.Client() as http: + response = http.post( + server.url.copy_with(path="/echo_body"), stream=True, data=b"*" * 409600 + ) + + assert response.status_code == 200 + with pytest.raises(RemoteProtocolError): + response.read() + + def test_stream_response(server): with httpx.Client() as http: response = http.get(server.url, stream=True) - assert response.status_code == 200 - content = response.read() - assert content == b"Hello, world!" + assert response.status_code == 200 + content = response.read() + assert content == b"Hello, world!" def test_stream_iterator(server): with httpx.Client() as http: response = http.get(server.url, stream=True) - assert response.status_code == 200 - body = b"" - for chunk in response.stream(): - body += chunk - assert body == b"Hello, world!" + assert response.status_code == 200 + body = b"" + for chunk in response.stream(): + body += chunk + assert body == b"Hello, world!" def test_raw_iterator(server): with httpx.Client() as http: response = http.get(server.url, stream=True) - assert response.status_code == 200 - body = b"" - for chunk in response.raw(): - body += chunk - assert body == b"Hello, world!" - response.close() # TODO: should Response be available as context managers? + assert response.status_code == 200 + body = b"" + for chunk in response.raw(): + body += chunk + assert body == b"Hello, world!" + response.close() # TODO: should Response be available as context managers? def test_raise_for_status(server): diff --git a/tests/dispatch/test_connection_pools.py b/tests/dispatch/test_connection_pools.py index 276580875f..5848a91922 100644 --- a/tests/dispatch/test_connection_pools.py +++ b/tests/dispatch/test_connection_pools.py @@ -112,6 +112,17 @@ async def test_standard_response_close(server, backend): assert len(http.keepalive_connections) == 1 +async def test_close_connection_on_context_manager_exit(server, backend): + """ + After block with all connections should be closed. + """ + async with httpx.ConnectionPool(backend=backend) as http: + await http.request("GET", server.url) + + assert len(http.active_connections) == 0 + assert len(http.keepalive_connections) == 0 + + async def test_premature_response_close(server, backend): """ A premature close should close the connection.