tests/test_cli.py

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

import pytest

from pytube import Caption
from pytube import CaptionQuery
from pytube import cli
from pytube import StreamQuery
from pytube.exceptions import PytubeError

parse_args = cli._parse_args


@mock.patch("pytube.cli._parse_args")
def test_main_invalid_url(_parse_args):  # noqa: PT019
    parser = argparse.ArgumentParser()
    args = parse_args(parser, ["crikey",],)
    _parse_args.return_value = args
    with pytest.raises(SystemExit):
        cli.main()


@mock.patch("pytube.cli.display_streams")
@mock.patch("pytube.cli.YouTube")
def test_download_when_itag_not_found(youtube, display_streams):
    # Given
    youtube.streams = mock.Mock()
    youtube.streams.get_by_itag.return_value = None
    # When
    with pytest.raises(SystemExit):
        cli.download_by_itag(youtube, 123)
    # Then
    youtube.streams.get_by_itag.assert_called_with(123)
    display_streams.assert_called_with(youtube)


@mock.patch("pytube.cli.YouTube")
@mock.patch("pytube.Stream")
def test_download_when_itag_is_found(youtube, stream):
    stream.itag = 123
    stream.exists_at_path.return_value = False
    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):
    # Given
    stream.itag = 123
    stream.__repr__ = MagicMock(return_value="")
    youtube.streams = StreamQuery([stream])
    # When
    cli.display_streams(youtube)
    # Then
    stream.__repr__.assert_called()


@mock.patch("pytube.cli._print_available_captions")
@mock.patch("pytube.cli.YouTube")
def test_download_caption_with_none(youtube, print_available):
    # Given
    caption = Caption(
        {"url": "url1", "name": {"simpleText": "name1"}, "languageCode": "en", "vssId": ".en"}
    )
    youtube.captions = CaptionQuery([caption])
    # When
    cli.download_caption(youtube, None)
    # Then
    print_available.assert_called_with(youtube.captions)


@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", "vssId": ".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._print_available_captions")
@mock.patch("pytube.cli.YouTube")
def test_download_caption_with_lang_not_found(youtube, print_available):
    # Given
    caption = Caption(
        {"url": "url1", "name": {"simpleText": "name1"}, "languageCode": "en", "vssId": ".en"}
    )
    youtube.captions = CaptionQuery([caption])
    # When
    cli.download_caption(youtube, "blah")
    # Then
    print_available.assert_called_with(youtube.captions)


def test_print_available_captions(capsys):
    # Given
    caption1 = Caption(
        {"url": "url1", "name": {"simpleText": "name1"}, "languageCode": "en", "vssId": ".en"}
    )
    caption2 = Caption(
        {"url": "url2", "name": {"simpleText": "name2"}, "languageCode": "fr", "vssId": ".fr"}
    )
    query = CaptionQuery([caption1, caption2])
    # When
    cli._print_available_captions(query)
    # Then
    captured = capsys.readouterr()
    assert captured.out == "Available caption codes are: en, fr\n"


def test_display_progress_bar(capsys):
    cli.display_progress_bar(bytes_received=25, filesize=100, scale=0.55)
    out, _ = capsys.readouterr()
    assert "25.0%" in out


@mock.patch("pytube.Stream")
def test_on_progress(stream):
    stream.filesize = 10
    cli.display_progress_bar = MagicMock()
    cli.on_progress(stream, "", 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",
            "-vvv",
        ],
    )
    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
    assert args.verbosity == 3


@mock.patch("pytube.cli.setup_logger", return_value=None)
def test_main_logging_setup(setup_logger):
    # Given
    parser = argparse.ArgumentParser()
    args = parse_args(parser, ["http://fakeurl", "-v"])
    cli._parse_args = MagicMock(return_value=args)
    # When
    with pytest.raises(SystemExit):
        cli.main()
    # Then
    setup_logger.assert_called_with(40, log_filename=None)


@mock.patch("pytube.cli.YouTube", 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", 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", 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", return_value=None)
def test_main_download_caption(youtube):
    parser = argparse.ArgumentParser()
    args = parse_args(parser, ["http://youtube.com/watch?v=9bZkp7q19f0", "-c", "en"])
    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", return_value=None)
@mock.patch("pytube.cli.download_by_resolution")
def test_download_by_resolution_flag(youtube, download_by_resolution):
    parser = argparse.ArgumentParser()
    args = parse_args(
        parser, ["http://youtube.com/watch?v=9bZkp7q19f0", "-r", "320p"]
    )
    cli._parse_args = MagicMock(return_value=args)
    cli.main()
    youtube.assert_called()
    download_by_resolution.assert_called()


@mock.patch("pytube.cli.YouTube")
@mock.patch("pytube.cli.Playlist")
@mock.patch("pytube.cli._perform_args_on_youtube")
def test_download_with_playlist(perform_args_on_youtube, playlist, youtube):
    # Given
    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)
    videos = [youtube]
    playlist_instance = playlist.return_value
    playlist_instance.videos = videos
    # When
    cli.main()
    # Then
    playlist.assert_called()
    perform_args_on_youtube.assert_called_with(youtube, args)


@mock.patch("pytube.cli.YouTube")
@mock.patch("pytube.cli.Playlist")
@mock.patch("pytube.cli._perform_args_on_youtube")
def test_download_with_playlist_video_error(
    perform_args_on_youtube, playlist, youtube, capsys
):
    # Given
    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)
    videos = [youtube]
    playlist_instance = playlist.return_value
    playlist_instance.videos = videos
    perform_args_on_youtube.side_effect = PytubeError()
    # When
    cli.main()
    # Then
    playlist.assert_called()
    captured = capsys.readouterr()
    assert "There was an error with video" in captured.out


@mock.patch("pytube.cli.YouTube")
@mock.patch("pytube.StreamQuery")
@mock.patch("pytube.Stream")
@mock.patch("pytube.cli._download")
def test_download_by_resolution(download, stream, stream_query, youtube):
    # Given
    stream_query.get_by_resolution.return_value = stream
    youtube.streams = stream_query
    # When
    cli.download_by_resolution(
        youtube=youtube, resolution="320p", target="test_target"
    )
    # Then
    download.assert_called_with(stream, target="test_target")


@mock.patch("pytube.cli.YouTube")
@mock.patch("pytube.StreamQuery")
@mock.patch("pytube.cli._download")
def test_download_by_resolution_not_exists(download, stream_query, youtube):
    stream_query.get_by_resolution.return_value = None
    youtube.streams = stream_query
    with pytest.raises(SystemExit):
        cli.download_by_resolution(
            youtube=youtube, resolution="DOESNT EXIST", target="test_target"
        )
    download.assert_not_called()


@mock.patch("pytube.Stream")
def test_download_stream_file_exists(stream, capsys):
    # Given
    stream.exists_at_path.return_value = True
    # When
    cli._download(stream=stream)
    # Then
    captured = capsys.readouterr()
    assert "Already downloaded at" in captured.out
    stream.download.assert_not_called()


@mock.patch("pytube.cli.YouTube")
@mock.patch("pytube.cli.ffmpeg_process")
def test_perform_args_should_ffmpeg_process(ffmpeg_process, youtube):
    # Given
    parser = argparse.ArgumentParser()
    args = parse_args(
        parser, ["http://youtube.com/watch?v=9bZkp7q19f0", "-f", "best"]
    )
    cli._parse_args = MagicMock(return_value=args)
    # When
    cli._perform_args_on_youtube(youtube, args)
    # Then
    ffmpeg_process.assert_called_with(
        youtube=youtube, resolution="best", target=None
    )


@mock.patch("pytube.cli.YouTube")
@mock.patch("pytube.cli._ffmpeg_downloader")
def test_ffmpeg_process_best_should_download(  # noqa: PT019
    _ffmpeg_downloader, youtube
):
    # Given
    target = "/target"
    streams = MagicMock()
    youtube.streams = streams
    video_stream = MagicMock()
    streams.filter.return_value.order_by.return_value.last.return_value = (
        video_stream
    )
    audio_stream = MagicMock()
    streams.get_audio_only.return_value = audio_stream
    # When
    cli.ffmpeg_process(youtube, "best", target)
    # Then
    _ffmpeg_downloader.assert_called_with(
        audio_stream=audio_stream, video_stream=video_stream, target=target
    )


@mock.patch("pytube.cli.YouTube")
@mock.patch("pytube.cli._ffmpeg_downloader")
def test_ffmpeg_process_res_should_download(  # noqa: PT019
    _ffmpeg_downloader, youtube
):
    # Given
    target = "/target"
    streams = MagicMock()
    youtube.streams = streams
    video_stream = MagicMock()
    streams.filter.return_value.first.return_value = video_stream
    audio_stream = MagicMock()
    streams.get_audio_only.return_value = audio_stream
    # When
    cli.ffmpeg_process(youtube, "XYZp", target)
    # Then
    _ffmpeg_downloader.assert_called_with(
        audio_stream=audio_stream, video_stream=video_stream, target=target
    )


@mock.patch("pytube.cli.YouTube")
@mock.patch("pytube.cli._ffmpeg_downloader")
def test_ffmpeg_process_res_none_should_not_download(  # noqa: PT019
    _ffmpeg_downloader, youtube
):
    # Given
    target = "/target"
    streams = MagicMock()
    youtube.streams = streams
    streams.filter.return_value.first.return_value = None
    audio_stream = MagicMock()
    streams.get_audio_only.return_value = audio_stream
    # When
    with pytest.raises(SystemExit):
        cli.ffmpeg_process(youtube, "XYZp", target)
    # Then
    _ffmpeg_downloader.assert_not_called()


@mock.patch("pytube.cli.YouTube")
@mock.patch("pytube.cli._ffmpeg_downloader")
def test_ffmpeg_process_audio_none_should_fallback_download(  # noqa: PT019
    _ffmpeg_downloader, youtube
):
    # Given
    target = "/target"
    streams = MagicMock()
    youtube.streams = streams
    stream = MagicMock()
    streams.filter.return_value.order_by.return_value.last.return_value = stream
    streams.get_audio_only.return_value = None
    # When
    cli.ffmpeg_process(youtube, "best", target)
    # Then
    _ffmpeg_downloader.assert_called_with(
        audio_stream=stream, video_stream=stream, target=target
    )


@mock.patch("pytube.cli.YouTube")
@mock.patch("pytube.cli._ffmpeg_downloader")
def test_ffmpeg_process_audio_fallback_none_should_exit(  # noqa: PT019
    _ffmpeg_downloader, youtube
):
    # Given
    target = "/target"
    streams = MagicMock()
    youtube.streams = streams
    stream = MagicMock()
    streams.filter.return_value.order_by.return_value.last.side_effect = [
        stream,
        stream,
        None,
    ]
    streams.get_audio_only.return_value = None
    # When
    with pytest.raises(SystemExit):
        cli.ffmpeg_process(youtube, "best", target)
    # Then
    _ffmpeg_downloader.assert_not_called()


@mock.patch("pytube.cli.os.unlink", return_value=None)
@mock.patch("pytube.cli.subprocess.run", return_value=None)
@mock.patch("pytube.cli._download", return_value=None)
@mock.patch("pytube.cli._unique_name", return_value=None)
def test_ffmpeg_downloader(unique_name, download, run, unlink):
    # Given
    target = "target"
    audio_stream = MagicMock()
    video_stream = MagicMock()
    video_stream.id = "video_id"
    audio_stream.subtype = "audio_subtype"
    video_stream.subtype = "video_subtype"
    unique_name.side_effect = ["video_name", "audio_name"]

    # When
    cli._ffmpeg_downloader(
        audio_stream=audio_stream, video_stream=video_stream, target=target
    )
    # Then
    download.assert_called()
    run.assert_called_with(
        [
            "ffmpeg",
            "-i",
            os.path.join("target", "video_name.video_subtype"),
            "-i",
            os.path.join("target", "audio_name.audio_subtype"),
            "-codec",
            "copy",
            os.path.join("target", "safe_title.video_subtype"),
        ]
    )
    unlink.assert_called()


@mock.patch("pytube.cli.download_audio")
@mock.patch("pytube.cli.YouTube.__init__", return_value=None)
def test_download_audio_args(youtube, download_audio):
    # Given
    parser = argparse.ArgumentParser()
    args = parse_args(
        parser, ["http://youtube.com/watch?v=9bZkp7q19f0", "-a", "mp4"]
    )
    cli._parse_args = MagicMock(return_value=args)
    # When
    cli.main()
    # Then
    youtube.assert_called()
    download_audio.assert_called()


@mock.patch("pytube.cli._download")
@mock.patch("pytube.cli.YouTube")
def test_download_audio(youtube, download):
    # Given
    youtube_instance = youtube.return_value
    audio_stream = MagicMock()
    youtube_instance.streams.filter.return_value.order_by.return_value.last.return_value = (
        audio_stream
    )
    # When
    cli.download_audio(youtube_instance, "filetype", "target")
    # Then
    download.assert_called_with(audio_stream, target="target")


@mock.patch("pytube.cli._download")
@mock.patch("pytube.cli.YouTube")
def test_download_audio_none(youtube, download):
    # Given
    youtube_instance = youtube.return_value
    youtube_instance.streams.filter.return_value.order_by.return_value.last.return_value = (
        None
    )
    # When
    with pytest.raises(SystemExit):
        cli.download_audio(youtube_instance, "filetype", "target")
    # Then
    download.assert_not_called()


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


@mock.patch("pytube.cli.os.path.exists", return_value=False)
def test_unique_name(path_exists):
    assert (
        cli._unique_name("base", "subtype", "video", "target") == "base_video_0"
    )


@mock.patch("pytube.cli.os.path.exists")
def test_unique_name_counter(path_exists):
    path_exists.side_effect = [True, False]
    assert (
        cli._unique_name("base", "subtype", "video", "target") == "base_video_1"
    )