[3.1] Call on_chunk_sent when write_eof is called with a valid chunk (G…

…H-2911) (cherry picked from commit 6a6ab7d) Co-authored-by: Pau Freixes <[email protected]>
aio-libs · Apr 4, 2018 · dc4f4ee · dc4f4ee
1 parent 98c7491
commit dc4f4ee
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 0 deletions.
diff --git a/CHANGES/2909.bugfix b/CHANGES/2909.bugfix
@@ -0,0 +1 @@
+Call on_chunk_sent when write_eof takes as a param the last chunk
diff --git a/aiohttp/http_writer.py b/aiohttp/http_writer.py
@@ -122,6 +122,9 @@ async def write_eof(self, chunk=b''):
                     chunk = b'0\r\n\r\n'
 
         if chunk:
+            if self._on_chunk_sent is not None:
+                await self._on_chunk_sent(chunk)
+
             self._write(chunk)
 
         await self.drain()

diff --git a/tests/test_http_writer.py b/tests/test_http_writer.py
@@ -171,6 +171,18 @@ async def test_write_calls_callback(protocol, transport, loop):
     assert on_chunk_sent.call_args == mock.call(chunk)
 
 
+async def test_write_eof_calls_callback(protocol, transport, loop):
+    on_chunk_sent = make_mocked_coro()
+    msg = http.StreamWriter(
+        protocol, transport, loop,
+        on_chunk_sent=on_chunk_sent
+    )
+    chunk = b'1'
+    await msg.write_eof(chunk=chunk)
+    assert on_chunk_sent.called
+    assert on_chunk_sent.call_args == mock.call(chunk)
+
+
 async def test_write_to_closing_transport(protocol, transport, loop):
     msg = http.StreamWriter(protocol, transport, loop)
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1 @@
		Call on_chunk_sent when write_eof takes as a param the last chunk