File size: 5,024 Bytes
d68bda8 eec6f64 e7964fa a92dd5e e7964fa a92dd5e e7964fa 3e182e5 e7964fa 3e182e5 e7964fa 82321d6 9172cb2 9e52f7b 9172cb2 9e52f7b b2b9291 272841b eec6f64 b2b9291 272841b b2b9291 272841b b2b9291 cf801d3 d68bda8 b22e0c8 d68bda8 af9a117 9e555f0 b22e0c8 fc9aec5 d68bda8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 |
import gzip
import io
import json
import os
import pytest
from unittest import mock
from pytube import helpers
from pytube.exceptions import RegexMatchError
from pytube.helpers import cache, create_mock_html_json, deprecated, setup_logger
from pytube.helpers import target_directory, uniqueify
def test_regex_search_no_match():
with pytest.raises(RegexMatchError):
helpers.regex_search("^a$", "", group=0)
def test_regex_search():
assert helpers.regex_search("^a$", "a", group=0) == "a"
def test_safe_filename():
"""Unsafe characters get stripped from generated filename"""
assert helpers.safe_filename("abc1245$$") == "abc1245"
assert helpers.safe_filename("abc##") == "abc"
@mock.patch("warnings.warn")
def test_deprecated(warn):
@deprecated("oh no")
def deprecated_function():
return None
deprecated_function()
warn.assert_called_with(
"Call to deprecated function deprecated_function (oh no).",
category=DeprecationWarning,
stacklevel=2,
)
def test_cache():
call_count = 0
@cache
def cached_func(stuff):
nonlocal call_count
call_count += 1
return stuff
cached_func("hi")
cached_func("hi")
cached_func("bye")
cached_func("bye")
assert call_count == 2
@mock.patch("os.path.isabs", return_value=False)
@mock.patch("os.getcwd", return_value="/cwd")
@mock.patch("os.makedirs")
def test_target_directory_with_relative_path(_, __, makedirs): # noqa: PT019
assert target_directory("test") == os.path.join("/cwd", "test")
makedirs.assert_called()
@mock.patch("os.path.isabs", return_value=True)
@mock.patch("os.makedirs")
def test_target_directory_with_absolute_path(_, makedirs): # noqa: PT019
assert target_directory("/test") == "/test"
makedirs.assert_called()
@mock.patch("os.getcwd", return_value="/cwd")
@mock.patch("os.makedirs")
def test_target_directory_with_no_path(_, makedirs): # noqa: PT019
assert target_directory() == "/cwd"
makedirs.assert_called()
@mock.patch("pytube.helpers.logging")
def test_setup_logger(logging):
# Given
logger = logging.getLogger.return_value
# When
setup_logger(20)
# Then
logging.getLogger.assert_called_with("pytube")
logger.addHandler.assert_called()
logger.setLevel.assert_called_with(20)
@mock.patch('builtins.open', new_callable=mock.mock_open)
@mock.patch('pytube.request.urlopen')
def test_create_mock_html_json(mock_url_open, mock_open):
video_id = '2lAe1cqCOXo'
gzip_html_filename = 'yt-video-%s-html.json.gz' % video_id
# Get the pytube directory in order to navigate to /tests/mocks
pytube_dir_path = os.path.abspath(
os.path.join(
os.path.dirname(__file__),
os.path.pardir
)
)
pytube_mocks_path = os.path.join(pytube_dir_path, 'tests', 'mocks')
gzip_html_filepath = os.path.join(pytube_mocks_path, gzip_html_filename)
# Mock the responses to YouTube
mock_url_open_object = mock.Mock()
# Order is:
# 1. watch_html -- must have jsurl match
# 2. embed html
# 3. watch html
# 4. raw vid info
mock_url_open_object.read.side_effect = [
(b'yt.setConfig({"PLAYER_CONFIG":{"args":[]}});ytInitialData = {};ytInitialPlayerResponse = {};' # noqa: E501
b'"jsUrl":"/s/player/13371337/player_ias.vflset/en_US/base.js"'),
b'embed_html',
b'watch_html',
b'{\"responseContext\":{}}',
]
mock_url_open.return_value = mock_url_open_object
# Generate a json with sample html json
result_data = create_mock_html_json(video_id)
# Assert that a write was only made once
mock_open.assert_called_once_with(gzip_html_filepath, 'wb')
# The result data should look like this:
gzip_file = io.BytesIO()
with gzip.GzipFile(
filename=gzip_html_filename,
fileobj=gzip_file,
mode='wb'
) as f:
f.write(json.dumps(result_data).encode('utf-8'))
gzip_data = gzip_file.getvalue()
file_handle = mock_open.return_value.__enter__.return_value
# For some reason, write gets called multiple times, so we have to
# concatenate all the write calls to get the full data before we compare
# it to the BytesIO object value.
full_content = b''
for call in file_handle.write.call_args_list:
args, kwargs = call
full_content += b''.join(args)
# The file header includes time metadata, so *occasionally* a single
# byte will be off at the very beginning. In theory, this difference
# should only affect bytes 5-8 (or [4:8] because of zero-indexing),
# but I've excluded the 10-byte metadata header altogether from the
# check, just to be safe.
# Source: https://en.wikipedia.org/wiki/Gzip#File_format
assert gzip_data[10:] == full_content[10:]
def test_uniqueify():
non_unique_list = [1, 2, 3, 3, 4, 5]
expected = [1, 2, 3, 4, 5]
result = uniqueify(non_unique_list)
assert result == expected
|