tests/test_cli.py

# -*- coding: utf-8 -*-
import argparse
from unittest import mock
from unittest.mock import MagicMock, patch

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):
    youtube.streams = mock.Mock()
    youtube.streams.all.return_value = []
    youtube.streams.get_by_itag.return_value = None
    with pytest.raises(SystemExit):
        cli.download_by_itag(youtube, 123)
    youtube.streams.get_by_itag.assert_called_with(123)


@mock.patch("pytube.cli.YouTube")
@mock.patch("pytube.Stream")
def test_download_when_itag_is_found(youtube, stream):
    stream.itag = 123
    youtube.streams = StreamQuery([stream])
    with patch.object(
        youtube.streams, "get_by_itag", wraps=youtube.streams.get_by_itag
    ) as wrapped_itag:
        cli.download_by_itag(youtube, 123)
        wrapped_itag.assert_called_with(123)
    youtube.register_on_progress_callback.assert_called_with(cli.on_progress)
    stream.download.assert_called()


@mock.patch("pytube.cli.YouTube")
@mock.patch("pytube.Stream")
def test_display_stream(youtube, stream):
    stream.itag = 123
    stream.__repr__ = MagicMock(return_value="")
    youtube.streams = StreamQuery([stream])
    with patch.object(youtube.streams, "all", wraps=youtube.streams.all) as wrapped_all:
        cli.display_streams(youtube)
        wrapped_all.assert_called()
        stream.__repr__.assert_called()


@mock.patch("pytube.cli.YouTube")
def test_download_caption_with_none(youtube):
    caption = Caption(
        {"url": "url1", "name": {"simpleText": "name1"}, "languageCode": "en"}
    )
    youtube.captions = CaptionQuery([caption])
    with patch.object(
        youtube.captions, "all", wraps=youtube.captions.all
    ) as wrapped_all:
        cli.download_caption(youtube, None)
        wrapped_all.assert_called()


@mock.patch("pytube.cli.YouTube")
def test_download_caption_with_language_found(youtube):
    youtube.title = "video title"
    caption = Caption(
        {"url": "url1", "name": {"simpleText": "name1"}, "languageCode": "en"}
    )
    caption.download = MagicMock(return_value="file_path")
    youtube.captions = CaptionQuery([caption])
    cli.download_caption(youtube, "en")
    caption.download.assert_called_with(title="video title", output_path=None)


@mock.patch("pytube.cli.YouTube")
def test_download_caption_with_language_not_found(youtube):
    caption = Caption(
        {"url": "url1", "name": {"simpleText": "name1"}, "languageCode": "en"}
    )
    youtube.captions = CaptionQuery([caption])
    with patch.object(
        youtube.captions, "all", wraps=youtube.captions.all
    ) as wrapped_all:
        cli.download_caption(youtube, "blah")
        wrapped_all.assert_called()


@mock.patch("pytube.Stream")
@mock.patch("io.BufferedWriter")
def test_on_progress(stream, writer):
    stream.filesize = 10
    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, ["http://youtube.com/watch?v=9bZkp7q19f0"])
    assert args.url == "http://youtube.com/watch?v=9bZkp7q19f0"
    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,
        [
            "http://youtube.com/watch?v=9bZkp7q19f0",
            "--build-playback-report",
            "-c",
            "en",
            "-l",
            "--itag=10",
        ],
    )
    assert args.url == "http://youtube.com/watch?v=9bZkp7q19f0"
    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_by_itag(youtube):
    parser = argparse.ArgumentParser()
    args = parse_args(parser, ["http://youtube.com/watch?v=9bZkp7q19f0", "--itag=10"])
    cli._parse_args = MagicMock(return_value=args)
    cli.download_by_itag = MagicMock()
    cli.main()
    youtube.assert_called()
    cli.download_by_itag.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, ["http://youtube.com/watch?v=9bZkp7q19f0", "--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, ["http://youtube.com/watch?v=9bZkp7q19f0", "-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, ["http://youtube.com/watch?v=9bZkp7q19f0", "-c"])
    cli._parse_args = MagicMock(return_value=args)
    cli.download_caption = MagicMock()
    cli.main()
    youtube.assert_called()
    cli.download_caption.assert_called()


@mock.patch("pytube.cli.YouTube.__init__", return_value=None)
def test_download_by_resolution(youtube):
    parser = argparse.ArgumentParser()
    args = parse_args(parser, ["http://youtube.com/watch?v=9bZkp7q19f0", "-r", "320p"])
    cli._parse_args = MagicMock(return_value=args)
    cli.download_by_resolution = MagicMock()
    cli.main()
    youtube.assert_called()
    cli.download_by_resolution.assert_called()


@mock.patch("pytube.cli.Playlist")
def test_download_with_playlist(playlist):
    cli.safe_filename = MagicMock(return_value="safe_title")
    parser = argparse.ArgumentParser()
    args = parse_args(parser, ["https://www.youtube.com/playlist?list=PLyn"])
    cli._parse_args = MagicMock(return_value=args)
    cli.download_by_resolution = MagicMock()
    cli.main()
    playlist.assert_called()