Merge pull request #21 from pytest-dev/chunked-encoding/1/dev

diazona · web-flow · commit 387eb4a9e2b9 · 2021-12-30T19:03:24.000-08:00
Support chunked encoding
diff --git a/README.rst b/README.rst
@@ -73,14 +73,25 @@ poking around in the code itself.
     following attributes:
 
     * ``code`` - HTTP response code (int)
-    * ``content`` - content of next response (str)
+    * ``content`` - content of next response (str, bytes, or iterable of either)
     * ``headers`` - response headers (dict)
+    * ``chunked`` - whether to chunk-encode the response (enumeration)
 
-    Once these attribute are set, all subsequent requests will be answered with
+    Once these attributes are set, all subsequent requests will be answered with
     these values until they are changed or the server is stopped. A more
     convenient way to change these is ::
 
-        httpserver.serve_content(content=None, code=200, headers=None)
+        httpserver.serve_content(content=None, code=200, headers=None, chunked=pytest_localserver.http.Chunked.NO)
+
+    The ``chunked`` atribute or parameter can be set to
+
+    * ``Chunked.YES``, telling the server to always apply chunk encoding
+    * ``Chunked.NO``, telling the server to never apply chunk encoding
+    * ``Chunked.AUTO``, telling the server to apply chunk encoding only if
+        the ``Transfer-Encoding`` header includes ``chunked``
+
+    If chunk encoding is applied, each str or bytes in ``content`` becomes one
+    chunk in the response.
 
     The server address can be found in property
 
diff --git a/pytest_localserver/http.py b/pytest_localserver/http.py
@@ -3,10 +3,13 @@
 # This program is release under the MIT license. You can find the full text of
 # the license in the LICENSE file.
 
+import enum
+import itertools
 import json
 import sys
 import threading
 
+from werkzeug.datastructures import Headers
 from werkzeug.serving import make_server
 from werkzeug.wrappers import Response, Request
 
@@ -39,6 +42,21 @@ def url(self):
         return '%s://%s:%i' % (proto, host, port)
 
 
+class Chunked(enum.Enum):
+    NO = False
+    YES = True
+    AUTO = None
+
+    def __bool__(self):
+        return bool(self.value)
+
+
+def _encode_chunk(chunk, charset):
+    if isinstance(chunk, str):
+        chunk = chunk.encode(charset)
+    return '{0:x}'.format(len(chunk)).encode(charset) + b'\r\n' + chunk + b'\r\n'
+
+
 class ContentServer(WSGIServer):
 
     """
@@ -67,6 +85,7 @@ def __init__(self, host='127.0.0.1', port=0, ssl_context=None):
         self.show_post_vars = False
         self.compress = None
         self.requests = []
+        self.chunked = Chunked.NO
 
     def __call__(self, environ, start_response):
         """
@@ -80,7 +99,21 @@ def __call__(self, environ, start_response):
         else:
             content = self.content
 
-        response = Response(status=self.code)
+        if (
+            self.chunked == Chunked.YES
+            or (self.chunked == Chunked.AUTO and 'chunked' in self.headers.get('Transfer-encoding', ''))
+        ):
+            # If the code below ever changes to allow setting the charset of
+            # the Response object, the charset used here should also be changed
+            # to match. But until that happens, use UTF-8 since it is Werkzeug's
+            # default.
+            charset = 'utf-8'
+            if isinstance(content, (str, bytes)):
+                content = (_encode_chunk(content, charset), '0\r\n\r\n')
+            else:
+                content = itertools.chain((_encode_chunk(item, charset) for item in content), ['0\r\n\r\n'])
+
+        response = Response(response=content, status=self.code)
         response.headers.clear()
         response.headers.extend(self.headers)
 
@@ -89,21 +122,34 @@ def __call__(self, environ, start_response):
         #     content = gzip.compress(content.encode('utf-8'))
         #     response.content_encoding = 'gzip'
 
-        response.data = content
         return response(environ, start_response)
 
-    def serve_content(self, content, code=200, headers=None):
+    def serve_content(self, content, code=200, headers=None, chunked=Chunked.NO):
         """
         Serves string content (with specified HTTP error code) as response to
         all subsequent request.
 
         :param content: content to be displayed
         :param code: HTTP status code
         :param headers: HTTP headers to be returned
+        :param chunked: whether to apply chunked transfer encoding to the content
         """
-        self.content, self.code = (content, code)
+        if not isinstance(content, (str, bytes, list, tuple)):
+            # If content is an iterable which is not known to be a string,
+            # bytes, or sequence, it might be something that can only be iterated
+            # through once, in which case we need to cache it so it can be reused
+            # to handle multiple requests.
+            try:
+                content = tuple(iter(content))
+            except TypeError:
+                # this probably means that content is not iterable, so just go
+                # ahead in case it's some type that Response knows how to handle
+                pass
+        self.content = content
+        self.code = code
+        self.chunked = chunked
         if headers:
-            self.headers = headers
+            self.headers = Headers(headers)
 
 
 if __name__ == '__main__':  # pragma: no cover
diff --git a/tests/test_http.py b/tests/test_http.py
@@ -1,3 +1,5 @@
+import itertools
+import pytest
 import requests
 
 from pytest_localserver import http, plugin
@@ -8,6 +10,9 @@
 httpserver = plugin.httpserver
 
 
+transfer_encoded = pytest.mark.parametrize('transfer_encoding_header', ['Transfer-encoding', 'Transfer-Encoding', 'transfer-encoding', 'TRANSFER-ENCODING'])
+
+
 def test_httpserver_funcarg(httpserver):
     assert isinstance(httpserver, http.ContentServer)
     assert httpserver.is_alive()
@@ -76,3 +81,206 @@ def test_HEAD_request(httpserver):
 #     resp = requests.post(httpserver.url, data={'data': 'value'}, headers=headers)
 #     assert resp.json() == {'data': 'value'}
 #     assert resp.status_code == 200
+
+
+@pytest.mark.parametrize('chunked_flag', [http.Chunked.YES, http.Chunked.AUTO, http.Chunked.NO])
+def test_chunked_attribute_without_header(httpserver, chunked_flag):
+    """
+    Test that passing the chunked attribute to serve_content() properly sets
+    the chunked property of the server.
+    """
+    httpserver.serve_content(
+        ('TEST!', 'test'),
+        headers={'Content-type': 'text/plain'},
+        chunked=chunked_flag
+    )
+    assert httpserver.chunked == chunked_flag
+
+
+@pytest.mark.parametrize('chunked_flag', [http.Chunked.YES, http.Chunked.AUTO, http.Chunked.NO])
+def test_chunked_attribute_with_header(httpserver, chunked_flag):
+    """
+    Test that passing the chunked attribute to serve_content() properly sets
+    the chunked property of the server even when the transfer-encoding header is
+    also set.
+    """
+    httpserver.serve_content(
+        ('TEST!', 'test'),
+        headers={'Content-type': 'text/plain', 'Transfer-encoding': 'chunked'},
+        chunked=chunked_flag
+    )
+    assert httpserver.chunked == chunked_flag
+
+
+@transfer_encoded
+@pytest.mark.parametrize('chunked_flag', [http.Chunked.YES, http.Chunked.AUTO])
+def test_GET_request_chunked_parameter(httpserver, transfer_encoding_header, chunked_flag):
+    """
+    Test that passing YES or AUTO as the chunked parameter to serve_content()
+    causes the response to be sent using chunking when the Transfer-encoding
+    header is also set.
+    """
+    httpserver.serve_content(
+        ('TEST!', 'test'),
+        headers={'Content-type': 'text/plain', transfer_encoding_header: 'chunked'},
+        chunked=chunked_flag
+    )
+    resp = requests.get(httpserver.url, headers={'User-Agent': 'Test method'})
+    assert resp.text == 'TEST!test'
+    assert resp.status_code == 200
+    assert 'text/plain' in resp.headers['Content-type']
+    assert 'chunked' in resp.headers['Transfer-encoding']
+
+
+@transfer_encoded
+@pytest.mark.parametrize('chunked_flag', [http.Chunked.YES, http.Chunked.AUTO])
+def test_GET_request_chunked_attribute(httpserver, transfer_encoding_header, chunked_flag):
+    """
+    Test that setting the chunked attribute of httpserver to YES or AUTO
+    causes the response to be sent using chunking when the Transfer-encoding
+    header is also set.
+    """
+    httpserver.serve_content(
+        ('TEST!', 'test'),
+        headers={'Content-type': 'text/plain', transfer_encoding_header: 'chunked'}
+    )
+    httpserver.chunked = chunked_flag
+    resp = requests.get(httpserver.url, headers={'User-Agent': 'Test method'})
+    assert resp.text == 'TEST!test'
+    assert resp.status_code == 200
+    assert 'text/plain' in resp.headers['Content-type']
+    assert 'chunked' in resp.headers['Transfer-encoding']
+
+
+@transfer_encoded
+def test_GET_request_not_chunked(httpserver, transfer_encoding_header):
+    """
+    Test that setting the chunked attribute of httpserver to NO causes
+    the response not to be sent using chunking even if the Transfer-encoding
+    header is set.
+    """
+    httpserver.serve_content(
+        ('TEST!', 'test'),
+        headers={'Content-type': 'text/plain', transfer_encoding_header: 'chunked'},
+        chunked=http.Chunked.NO
+    )
+    with pytest.raises(requests.exceptions.ChunkedEncodingError):
+        resp = requests.get(httpserver.url, headers={'User-Agent': 'Test method'})
+
+
+@pytest.mark.parametrize('chunked_flag', [http.Chunked.NO, http.Chunked.AUTO])
+def test_GET_request_chunked_parameter_no_header(httpserver, chunked_flag):
+    """
+    Test that passing NO or AUTO as the chunked parameter to serve_content()
+    causes the response not to be sent using chunking when the Transfer-encoding
+    header is not set.
+    """
+    httpserver.serve_content(
+        ('TEST!', 'test'),
+        headers={'Content-type': 'text/plain'},
+        chunked=chunked_flag
+    )
+    resp = requests.get(httpserver.url, headers={'User-Agent': 'Test method'})
+    assert resp.text == 'TEST!test'
+    assert resp.status_code == 200
+    assert 'text/plain' in resp.headers['Content-type']
+    assert 'Transfer-encoding' not in resp.headers
+
+
+@pytest.mark.parametrize('chunked_flag', [http.Chunked.NO, http.Chunked.AUTO])
+def test_GET_request_chunked_attribute_no_header(httpserver, chunked_flag):
+    """
+    Test that setting the chunked attribute of httpserver to NO or AUTO
+    causes the response not to be sent using chunking when the Transfer-encoding
+    header is not set.
+    """
+    httpserver.serve_content(
+        ('TEST!', 'test'),
+        headers={'Content-type': 'text/plain'}
+    )
+    httpserver.chunked = chunked_flag
+    resp = requests.get(httpserver.url, headers={'User-Agent': 'Test method'})
+    assert resp.text == 'TEST!test'
+    assert resp.status_code == 200
+    assert 'text/plain' in resp.headers['Content-type']
+    assert 'Transfer-encoding' not in resp.headers
+
+
+def test_GET_request_chunked_no_header(httpserver):
+    """
+    Test that setting the chunked attribute of httpserver to YES causes
+    the response to be sent using chunking even if the Transfer-encoding
+    header is not set.
+    """
+    httpserver.serve_content(
+        ('TEST!', 'test'),
+        headers={'Content-type': 'text/plain'},
+        chunked=http.Chunked.YES
+    )
+    resp = requests.get(httpserver.url, headers={'User-Agent': 'Test method'})
+    # Without the Transfer-encoding header set, requests does not undo the chunk
+    # encoding so it comes through as "raw" chunks
+    assert resp.text == '5\r\nTEST!\r\n4\r\ntest\r\n0\r\n\r\n'
+
+
+def _format_chunk(chunk):
+    r = repr(chunk)
+    if len(r) <= 40:
+        return r
+    else:
+        return r[:13] + '...' + r[-14:] + ' (length {0})'.format(len(chunk))
+
+
+def _compare_chunks(expected, actual):
+    __tracebackhide__ = True
+    if expected != actual:
+        message = [_format_chunk(expected) + ' != ' + _format_chunk(actual)]
+        if type(expected) == type(actual):
+            for i, (e, a) in enumerate(itertools.zip_longest(expected, actual, fillvalue='<end>')):
+                if e != a:
+                    message += [
+                        '  Chunks differ at index {}:'.format(i),
+                        '    Expected: ' + (repr(expected[i:i+5]) + '...' if e != '<end>' else '<end>'),
+                        '    Found:    ' + (repr(actual[i:i+5]) + '...' if a != '<end>' else '<end>')
+                    ]
+                    break
+        pytest.fail('\n'.join(message))
+
+
+@pytest.mark.parametrize('chunk_size', [400, 499, 500, 512, 750, 1024, 4096, 8192])
+def test_GET_request_large_chunks(httpserver, chunk_size):
+    """
+    Test that a response with large chunks comes through correctly
+    """
+    body = b'0123456789abcdef' * 1024  # 16 kb total
+    # Split body into fixed-size chunks, from https://stackoverflow.com/a/18854817/56541
+    chunks = [body[0 + i:chunk_size + i] for i in range(0, len(body), chunk_size)]
+    httpserver.serve_content(
+        chunks,
+        headers={'Content-type': 'text/plain', 'Transfer-encoding': 'chunked'},
+        chunked=http.Chunked.YES
+    )
+    resp = requests.get(httpserver.url, headers={'User-Agent': 'Test method'}, stream=True)
+    assert resp.status_code == 200
+    text = b''
+    for original_chunk, received_chunk in itertools.zip_longest(chunks, resp.iter_content(chunk_size=None)):
+        _compare_chunks(original_chunk, received_chunk)
+        text += received_chunk
+    assert text == body
+    assert 'chunked' in resp.headers['Transfer-encoding']
+
+
+@pytest.mark.parametrize('chunked_flag', [http.Chunked.YES, http.Chunked.AUTO])
+def test_GET_request_chunked_no_content_length(httpserver, chunked_flag):
+    """
+    Test that a chunked response does not include a Content-length header
+    """
+    httpserver.serve_content(
+        ('TEST!', 'test'),
+        headers={'Content-type': 'text/plain', 'Transfer-encoding': 'chunked'},
+        chunked=chunked_flag
+    )
+    resp = requests.get(httpserver.url, headers={'User-Agent': 'Test method'})
+    assert resp.status_code == 200
+    assert 'Transfer-encoding' in resp.headers
+    assert 'Content-length' not in resp.headers
