Skip to content

Commit

Permalink
Preserve chunked argument in urlopen() retry (#1715)
Browse files Browse the repository at this point in the history
  • Loading branch information
lmvlmv authored and sethmlarson committed Oct 29, 2019
1 parent 79e81f9 commit 31db282
Show file tree
Hide file tree
Showing 3 changed files with 31 additions and 3 deletions.
9 changes: 7 additions & 2 deletions dummyserver/testcase.py
Expand Up @@ -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):
Expand Down
1 change: 1 addition & 0 deletions src/urllib3/connectionpool.py
Expand Up @@ -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
)
Expand Down
24 changes: 23 additions & 1 deletion 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):
Expand Down Expand Up @@ -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

0 comments on commit 31db282

Please sign in to comment.