diff --git a/adafruit_requests.py b/adafruit_requests.py index 483bb80..7d7e86d 100644 --- a/adafruit_requests.py +++ b/adafruit_requests.py @@ -220,7 +220,7 @@ def _readinto(self, buf): # Consume trailing \r\n for chunks 2+ if self._remaining == 0: self._throw_away(2) - chunk_header = self._readto(b";", b"\r\n") + chunk_header = self._readto(b"\r\n").split(b";", 1)[0] http_chunk_size = int(bytes(chunk_header), 16) if http_chunk_size == 0: self._chunked = False @@ -261,7 +261,7 @@ def close(self): self._throw_away(self._remaining) elif self._chunked: while True: - chunk_header = self._readto(b";", b"\r\n") + chunk_header = self._readto(b"\r\n").split(b";", 1)[0] chunk_size = int(bytes(chunk_header), 16) if chunk_size == 0: break diff --git a/tests/chunk_test.py b/tests/chunk_test.py index ddbd048..0982652 100644 --- a/tests/chunk_test.py +++ b/tests/chunk_test.py @@ -9,7 +9,7 @@ headers = b"HTTP/1.0 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n" -def _chunk(response, split): +def _chunk(response, split, extra=b""): i = 0 chunked = b"" while i < len(response): @@ -19,7 +19,11 @@ def _chunk(response, split): chunk_size = remaining new_i = i + chunk_size chunked += ( - hex(chunk_size)[2:].encode("ascii") + b"\r\n" + response[i:new_i] + b"\r\n" + hex(chunk_size)[2:].encode("ascii") + + extra + + b"\r\n" + + response[i:new_i] + + b"\r\n" ) i = new_i # The final chunk is zero length. @@ -27,10 +31,10 @@ def _chunk(response, split): return chunked -def test_get_text(): +def do_test_get_text(extra=b""): pool = mocket.MocketPool() pool.getaddrinfo.return_value = ((None, None, None, None, (ip, 80)),) - c = _chunk(text, 33) + c = _chunk(text, 33, extra) print(c) sock = mocket.Mocket(headers + c) pool.socket.return_value = sock @@ -54,11 +58,19 @@ def test_get_text(): assert r.text == str(text, "utf-8") -def test_close_flush(): +def test_get_text(): + do_test_get_text() + + +def test_get_text_extra(): + do_test_get_text(b";blahblah; blah") + + +def do_test_close_flush(extra=b""): """Test that a chunked response can be closed even when the request contents were not accessed.""" pool = mocket.MocketPool() pool.getaddrinfo.return_value = ((None, None, None, None, (ip, 80)),) - c = _chunk(text, 33) + c = _chunk(text, 33, extra) print(c) sock = mocket.Mocket(headers + c) pool.socket.return_value = sock @@ -81,3 +93,11 @@ def test_close_flush(): ) r.close() + + +def test_close_flush(): + do_test_close_flush() + + +def test_close_flush_extra(): + do_test_close_flush(b";blahblah; blah")