cli: update progress output after stream has ended

bastimeyer · bastimeyer · commit 870fda060d93 · 2024-03-12T19:56:46.000+01:00
diff --git a/src/streamlink_cli/utils/progress.py b/src/streamlink_cli/utils/progress.py
@@ -251,9 +251,10 @@ def write(self, chunk: bytes):
     def run(self):
         self.started = time()
         try:
-            while not self._wait.wait(self.interval):  # pragma: no cover
+            while not self._wait.wait(self.interval):
                 self.update()
         finally:
+            self.update()
             self.print_end()
 
     def update(self):
diff --git a/tests/cli/utils/test_progress.py b/tests/cli/utils/test_progress.py
@@ -8,6 +8,7 @@
 import pytest
 
 from streamlink_cli.utils.progress import Progress, ProgressFormatter
+from tests.testutils.handshake import Handshake
 
 
 class TestProgressFormatter:
@@ -294,3 +295,43 @@ def test_download_speed(self):
             progress.update()
             assert output_write.call_args_list[-1] \
                    == call("\r[download] Written 21.00 KiB (6s @ 6.00 KiB/s)              ")
+
+    def test_update(self):
+        handshake = Handshake()
+
+        class _Progress(Progress):
+            def update(self):
+                with handshake():
+                    return super().update()
+
+        stream = StringIO()
+        thread = _Progress(stream=stream, path=PurePath())
+        # override the thread's polling time after initializing the deque of the rolling average download speed:
+        # the interval constructor keyword is used to set the deque size
+        thread.interval = 0
+        thread.start()
+
+        # first tick
+        assert handshake.wait_ready(1)
+        thread.write(b"123")
+        assert handshake.step(1)
+        assert stream.getvalue().split("\r")[-1].startswith("[download] Written 3 bytes")
+
+        # second tick
+        assert handshake.wait_ready(1)
+        thread.write(b"465")
+        assert handshake.step(1)
+        assert stream.getvalue().split("\r")[-1].startswith("[download] Written 6 bytes")
+
+        # close progress thread
+        assert handshake.wait_ready(1)
+        thread.close()
+        assert handshake.step(1)
+        assert stream.getvalue().split("\r")[-1].startswith("[download] Written 6 bytes")
+
+        # write data right after closing the thread, but before it has halted
+        thread.write(b"789")
+        handshake.go()
+        thread.join(1)
+        assert not thread.is_alive()
+        assert stream.getvalue().split("\r")[-1].startswith("[download] Written 9 bytes")
