Merge pull request #356 from PrimozGodec/last-chunk-fix

Fix: stream end when length of last chunk equal to MAX_CHUNK

Author: Lukasa

HISTORY.rst

@@ -4,6 +4,10 @@ Release History
 dev
 ---
 
+*Bugfixes*
+
+- Stream end flag when length of last chunk equal to MAX_CHUNK
+
 v0.7.0 (2016-09-27)
 -------------------
 

hyper/http20/stream.py

@@ -122,8 +122,18 @@ def file_iterator(fobj):
             chunks = (data[i:i+MAX_CHUNK]
                       for i in range(0, len(data), MAX_CHUNK))
 
-        for chunk in chunks:
-            self._send_chunk(chunk, final)
+        # since we need to know when we have a last package we need to know
+        # if there is another package in advance
+        cur_chunk = None
+        try:
+            cur_chunk = next(chunks)
+            while True:
+                next_chunk = next(chunks)
+                self._send_chunk(cur_chunk, False)
+                cur_chunk = next_chunk
+        except StopIteration:
+            if cur_chunk is not None:  # cur_chunk none when no chunks to send
+                self._send_chunk(cur_chunk, final)
 
     def _read(self, amt=None):
         """
@@ -323,19 +333,12 @@ def _send_chunk(self, data, final):
         while len(data) > self._out_flow_control_window:
             self._recv_cb()
 
-        # If the length of the data is less than MAX_CHUNK, we're probably
-        # at the end of the file. If this is the end of the data, mark it
-        # as END_STREAM.
-        end_stream = False
-        if len(data) < MAX_CHUNK and final:
-            end_stream = True
-
         # Send the frame and decrement the flow control window.
         with self._conn as conn:
             conn.send_data(
-                stream_id=self.stream_id, data=data, end_stream=end_stream
+                stream_id=self.stream_id, data=data, end_stream=final
             )
         self._send_outstanding_data()
 
-        if end_stream:
+        if final:
             self.local_closed = True

test/test_hyper.py

@@ -211,6 +211,61 @@ def data_callback(chunk, **kwargs):
         assert frames[1].data == b'hello there'
         assert frames[1].flags == set(['END_STREAM'])
 
+    def test_request_correctly_sent_max_chunk(self, frame_buffer):
+        """
+        Test that request correctly sent when data length multiple
+        max chunk. We check last chunk has a end flag and correct number
+        of chunks.
+        """
+        def data_callback(chunk, **kwargs):
+            frame_buffer.add_data(chunk)
+
+        # one chunk
+        c = HTTP20Connection('www.google.com')
+        c._sock = DummySocket()
+        c._send_cb = data_callback
+        c.putrequest('GET', '/')
+        c.endheaders(message_body=b'1'*1024, final=True)
+
+        frames = list(frame_buffer)
+        assert len(frames) == 2
+        assert isinstance(frames[1], DataFrame)
+        assert frames[1].flags == set(['END_STREAM'])
+
+        # two chunks
+        c = HTTP20Connection('www.google.com')
+        c._sock = DummySocket()
+        c._send_cb = data_callback
+        c.putrequest('GET', '/')
+        c.endheaders(message_body=b'1' * 2024, final=True)
+
+        frames = list(frame_buffer)
+        assert len(frames) == 3
+        assert isinstance(frames[1], DataFrame)
+        assert frames[2].flags == set(['END_STREAM'])
+
+        # two chunks with last chunk < 1024
+        c = HTTP20Connection('www.google.com')
+        c._sock = DummySocket()
+        c._send_cb = data_callback
+        c.putrequest('GET', '/')
+        c.endheaders(message_body=b'1' * 2000, final=True)
+
+        frames = list(frame_buffer)
+        assert len(frames) == 3
+        assert isinstance(frames[1], DataFrame)
+        assert frames[2].flags == set(['END_STREAM'])
+
+        # no chunks
+        c = HTTP20Connection('www.google.com')
+        c._sock = DummySocket()
+        c._send_cb = data_callback
+        c.putrequest('GET', '/')
+        c.endheaders(message_body=b'', final=True)
+
+        frames = list(frame_buffer)
+        assert len(frames) == 1
+
     def test_that_we_correctly_send_over_the_socket(self):
         sock = DummySocket()
         c = HTTP20Connection('www.google.com')