test for main cli and arguments parsing

pytube/cli.py

@@ -9,7 +9,7 @@ import logging
 import os
 import sys
 from io import BufferedWriter
-from typing import Tuple, Any, Optional
+from typing import Tuple, Any, Optional, List
 
 from pytube import __version__, CaptionQuery
 from pytube import YouTube
@@ -22,6 +22,28 @@ def main():
     """Command line application to download youtube videos."""
     # noinspection PyTypeChecker
     parser = argparse.ArgumentParser(description=main.__doc__)
+    args = _parse_args(parser)
+    logging.getLogger().setLevel(max(3 - args.verbosity, 0) * 10)
+
+    if not args.url:
+        parser.print_help()
+        sys.exit(1)
+
+    youtube = YouTube(args.url)
+
+    if args.list:
+        display_streams(youtube)
+    if args.build_playback_report:
+        build_playback_report(youtube)
+    if args.itag:
+        download(youtube=youtube, itag=args.itag)
+    if hasattr(args, "caption_code"):
+        download_caption(youtube=youtube, lang_code=args.caption_code)
+
+
+def _parse_args(
+    parser: argparse.ArgumentParser, args: Optional[List] = None
+) -> argparse.Namespace:
     parser.add_argument("url", help="The YouTube /watch url", nargs="?")
     parser.add_argument(
         "--version", action="version", version="%(prog)s " + __version__,
@@ -63,23 +85,7 @@ def main():
         ),
     )
 
-    args = parser.parse_args()
-    logging.getLogger().setLevel(max(3 - args.verbosity, 0) * 10)
-
-    if not args.url:
-        parser.print_help()
-        sys.exit(1)
-
-    youtube = YouTube(args.url)
-
-    if args.list:
-        display_streams(youtube)
-    if args.build_playback_report:
-        build_playback_report(youtube)
-    if args.itag:
-        download(youtube=youtube, itag=args.itag)
-    if hasattr(args, "caption_code"):
-        download_caption(youtube=youtube, lang_code=args.caption_code)
+    return parser.parse_args(args)
 
 
 def build_playback_report(youtube: YouTube) -> None:

tests/test_cli.py

@@ -1,4 +1,5 @@
 # -*- coding: utf-8 -*-
+import argparse
 from unittest import mock
 from unittest.mock import MagicMock, patch
 
@@ -6,6 +7,8 @@ import pytest
 
 from pytube import cli, StreamQuery, Caption, CaptionQuery
 
+parse_args = cli._parse_args
+
 
 @mock.patch("pytube.cli.YouTube")
 def test_download_when_itag_not_found(youtube):
@@ -88,3 +91,68 @@ def test_on_progress(stream, writer):
     cli.display_progress_bar = MagicMock()
     cli.on_progress(stream, "", writer, 7)
     cli.display_progress_bar.assert_called_once_with(3, 10)
+
+
+def test_parse_args_falsey():
+    parser = argparse.ArgumentParser()
+    args = cli._parse_args(parser, ["urlhere"])
+    assert args.url == "urlhere"
+    assert args.build_playback_report is False
+    assert args.itag is None
+    assert args.list is False
+    assert args.verbosity == 0
+
+
+def test_parse_args_truthy():
+    parser = argparse.ArgumentParser()
+    args = cli._parse_args(
+        parser, ["urlhere", "--build-playback-report", "-c", "en", "-l", "--itag=10"]
+    )
+    assert args.url == "urlhere"
+    assert args.build_playback_report is True
+    assert args.itag == 10
+    assert args.list is True
+
+
+@mock.patch("pytube.cli.YouTube.__init__", return_value=None)
+def test_main_download(youtube):
+    parser = argparse.ArgumentParser()
+    args = parse_args(parser, ["urlhere", "--itag=10"])
+    cli._parse_args = MagicMock(return_value=args)
+    cli.download = MagicMock()
+    cli.main()
+    youtube.assert_called()
+    cli.download.assert_called()
+
+
+@mock.patch("pytube.cli.YouTube.__init__", return_value=None)
+def test_main_build_playback_report(youtube):
+    parser = argparse.ArgumentParser()
+    args = parse_args(parser, ["urlhere", "--build-playback-report"])
+    cli._parse_args = MagicMock(return_value=args)
+    cli.build_playback_report = MagicMock()
+    cli.main()
+    youtube.assert_called()
+    cli.build_playback_report.assert_called()
+
+
+@mock.patch("pytube.cli.YouTube.__init__", return_value=None)
+def test_main_display_streams(youtube):
+    parser = argparse.ArgumentParser()
+    args = parse_args(parser, ["urlhere", "-l"])
+    cli._parse_args = MagicMock(return_value=args)
+    cli.display_streams = MagicMock()
+    cli.main()
+    youtube.assert_called()
+    cli.display_streams.assert_called()
+
+
+@mock.patch("pytube.cli.YouTube.__init__", return_value=None)
+def test_main_download_caption(youtube):
+    parser = argparse.ArgumentParser()
+    args = parse_args(parser, ["urlhere", "-c"])
+    cli._parse_args = MagicMock(return_value=args)
+    cli.download_caption = MagicMock()
+    cli.main()
+    youtube.assert_called()
+    cli.download_caption.assert_called()