more coverage

pytube/__main__.py

@@ -169,7 +169,7 @@ class YouTube(object):
             A callback function that takes ``stream``, ``chunk``,
             ``file_handle``, ``bytes_remaining`` as parameters.
         """
-        self._monostate['on_progress'] = func
+        self.stream_monostate['on_progress'] = func
 
     def register_on_complete_callback(self, func):
         """Register a download complete callback function post initialization.
@@ -177,4 +177,4 @@ class YouTube(object):
         :param callable func:
             A callback function that takes ``stream`` and  ``file_handle``.
         """
-        self._monostate['on_complete'] = func
+        self.stream_monostate['on_complete'] = func

pytube/streams.py

@@ -170,6 +170,7 @@ class Stream(object):
                 bytes_remaining -= len(chunk)
                 # send to the on_progress callback.
                 self.on_progress(chunk, fh, bytes_remaining)
+            self.on_complete(fh)
 
     def on_progress(self, chunk, file_handler, bytes_remaining):
         """On progress callback function.
@@ -202,7 +203,7 @@ class Stream(object):
         on_progress = self._monostate['on_progress']
         if on_progress:
             logger.debug('calling on_progress callback %s', on_progress)
-            on_progress(self, chunk, bytes_remaining)
+            on_progress(self, chunk, file_handler, bytes_remaining)
 
     def on_complete(self, file_handle):
         """On download complete handler function.

setup.cfg

@@ -4,7 +4,7 @@ tag = True
 current_version = 7.0.1
 parse = (?P<major>\d+)\.(?P<minor>\d+)\.(?P<patch>\d+)(\-(?P<release>[a-z]+))?
 serialize =
-	{major}.{minor}.{patch}
+    {major}.{minor}.{patch}
 
 [metadata]
 description-file = README.md
@@ -15,3 +15,5 @@ description-file = README.md
 
 [coverage:run]
 source = pytube
+omit =
+    pytube/compat.py

tests/test_extract.py

@@ -27,3 +27,9 @@ def test_info_url(gangnam_style):
         '253D9bZkp7q19f0&hl=en_US&t=%252C%2522t%2522%253A%25221%2522'
     )
     assert video_info_url == expected
+
+
+def test_js_url(gangnam_style):
+    expected = 'https://youtube.com/yts/jsbin/player-vflOdyxa4/en_US/base.js'
+    result = extract.js_url(gangnam_style.watch_html)
+    assert expected == result

tests/test_streams.py

@@ -4,6 +4,7 @@ import random
 import mock
 
 from pytube import request
+from pytube import Stream
 
 
 def test_filesize(gangnam_style, mocker):
@@ -38,3 +39,75 @@ def test_progressive_streams_return_includes_audio_track(gangnam_style):
 def test_progressive_streams_return_includes_video_track(gangnam_style):
     stream = gangnam_style.streams.filter(progressive=True).first()
     assert stream.includes_video_track
+
+
+def test_on_progress_hook(gangnam_style, mocker):
+    callback_fn = mock.MagicMock()
+    gangnam_style.register_on_progress_callback(callback_fn)
+
+    mocker.patch.object(request, 'get')
+    request.get.side_effect = [
+        {'Content-Length': '16384'},
+        {'Content-Length': '16384'},
+        iter([str(random.getrandbits(8 * 1024))]),
+    ]
+    with mock.patch('pytube.streams.open', mock.mock_open(), create=True):
+        stream = gangnam_style.streams.first()
+        stream.download()
+    assert callback_fn.called
+    args, _ = callback_fn.call_args
+    assert len(args) == 4
+    stream, _, _, _ = args
+    assert isinstance(stream, Stream)
+
+
+def test_on_complete_hook(gangnam_style, mocker):
+    callback_fn = mock.MagicMock()
+    gangnam_style.register_on_complete_callback(callback_fn)
+
+    mocker.patch.object(request, 'get')
+    request.get.side_effect = [
+        {'Content-Length': '16384'},
+        {'Content-Length': '16384'},
+        iter([str(random.getrandbits(8 * 1024))]),
+    ]
+    with mock.patch('pytube.streams.open', mock.mock_open(), create=True):
+        stream = gangnam_style.streams.first()
+        stream.download()
+    assert callback_fn.called
+
+
+def test_repr_for_audio_streams(gangnam_style):
+    stream = str(gangnam_style.streams.filter(only_audio=True).first())
+    expected = (
+        '<Stream: itag="140" mime_type="audio/mp4" abr="128kbps" '
+        'acodec="mp4a.40.2">'
+    )
+    assert stream == expected
+
+
+def test_repr_for_video_streams(gangnam_style):
+    stream = str(gangnam_style.streams.filter(only_video=True).first())
+    expected = (
+        '<Stream: itag="137" mime_type="video/mp4" res="1080p" '
+        'fps="30fps" vcodec="avc1.640028">'
+    )
+    assert stream == expected
+
+
+def test_repr_for_progressive_streams(gangnam_style):
+    stream = str(gangnam_style.streams.filter(progressive=True).first())
+    expected = (
+        '<Stream: itag="22" mime_type="video/mp4" res="720p" '
+        'fps="30fps" vcodec="avc1.64001F" acodec="mp4a.40.2">'
+    )
+    assert stream == expected
+
+
+def test_repr_for_adaptive_streams(gangnam_style):
+    stream = str(gangnam_style.streams.filter(adaptive=True).first())
+    expected = (
+        '<Stream: itag="137" mime_type="video/mp4" res="1080p" '
+        'fps="30fps" vcodec="avc1.640028">'
+    )
+    assert stream == expected