diff --git a/dummyserver/testcase.py b/dummyserver/testcase.py index bf83b6b432..07337d9a56 100644 --- a/dummyserver/testcase.py +++ b/dummyserver/testcase.py @@ -15,8 +15,13 @@ def consume_socket(sock, chunks=65536): - while not sock.recv(chunks).endswith(b"\r\n\r\n"): - pass + consumed = bytearray() + while True: + b = sock.recv(chunks) + consumed += b + if b.endswith(b"\r\n\r\n"): + break + return consumed class SocketDummyServerTestCase(object): diff --git a/src/urllib3/connectionpool.py b/src/urllib3/connectionpool.py index 58ce1cf827..26c4bd70ab 100644 --- a/src/urllib3/connectionpool.py +++ b/src/urllib3/connectionpool.py @@ -839,6 +839,7 @@ def drain_and_release_conn(response): timeout=timeout, pool_timeout=pool_timeout, release_conn=release_conn, + chunked=chunked, body_pos=body_pos, **response_kw ) diff --git a/test/with_dummyserver/test_chunked_transfer.py b/test/with_dummyserver/test_chunked_transfer.py index a7512cef69..a0807d5b66 100644 --- a/test/with_dummyserver/test_chunked_transfer.py +++ b/test/with_dummyserver/test_chunked_transfer.py @@ -1,7 +1,8 @@ # -*- coding: utf-8 -*- from urllib3 import HTTPConnectionPool -from dummyserver.testcase import SocketDummyServerTestCase +from urllib3.util.retry import Retry +from dummyserver.testcase import SocketDummyServerTestCase, consume_socket class TestChunkedTransfer(SocketDummyServerTestCase): @@ -100,3 +101,24 @@ def test_provides_default_host_header(self): host_headers = [x for x in header_lines if x.startswith(b"host")] assert len(host_headers) == 1 + + def test_preserve_chunked_on_retry(self): + self.chunked_requests = 0 + + def socket_handler(listener): + for _ in range(2): + sock = listener.accept()[0] + request = consume_socket(sock) + if b"Transfer-Encoding: chunked" in request.split(b"\r\n"): + self.chunked_requests += 1 + + sock.send(b"HTTP/1.1 404 Not Found\r\n\r\n") + sock.close() + + self._start_server(socket_handler) + with HTTPConnectionPool(self.host, self.port) as pool: + retries = Retry(total=1, raise_on_status=False, status_forcelist=[404]) + pool.urlopen( + "GET", "/", chunked=True, preload_content=False, retries=retries + ) + assert self.chunked_requests == 2