1from __future__ import unicode_literals
2
3import errno
4import io
5import hashlib
6import json
7import os.path
8import re
9import types
10import ssl
11import sys
12
13import yt_dlp.extractor
14from yt_dlp import YoutubeDL
15from yt_dlp.compat import (
16    compat_os_name,
17    compat_str,
18)
19from yt_dlp.utils import (
20    preferredencoding,
21    write_string,
22)
23
24
25if 'pytest' in sys.modules:
26    import pytest
27    is_download_test = pytest.mark.download
28else:
29    def is_download_test(testClass):
30        return testClass
31
32
33def get_params(override=None):
34    PARAMETERS_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)),
35                                   'parameters.json')
36    LOCAL_PARAMETERS_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)),
37                                         'local_parameters.json')
38    with io.open(PARAMETERS_FILE, encoding='utf-8') as pf:
39        parameters = json.load(pf)
40    if os.path.exists(LOCAL_PARAMETERS_FILE):
41        with io.open(LOCAL_PARAMETERS_FILE, encoding='utf-8') as pf:
42            parameters.update(json.load(pf))
43    if override:
44        parameters.update(override)
45    return parameters
46
47
48def try_rm(filename):
49    """ Remove a file if it exists """
50    try:
51        os.remove(filename)
52    except OSError as ose:
53        if ose.errno != errno.ENOENT:
54            raise
55
56
57def report_warning(message):
58    '''
59    Print the message to stderr, it will be prefixed with 'WARNING:'
60    If stderr is a tty file the 'WARNING:' will be colored
61    '''
62    if sys.stderr.isatty() and compat_os_name != 'nt':
63        _msg_header = '\033[0;33mWARNING:\033[0m'
64    else:
65        _msg_header = 'WARNING:'
66    output = '%s %s\n' % (_msg_header, message)
67    if 'b' in getattr(sys.stderr, 'mode', '') or sys.version_info[0] < 3:
68        output = output.encode(preferredencoding())
69    sys.stderr.write(output)
70
71
72class FakeYDL(YoutubeDL):
73    def __init__(self, override=None):
74        # Different instances of the downloader can't share the same dictionary
75        # some test set the "sublang" parameter, which would break the md5 checks.
76        params = get_params(override=override)
77        super(FakeYDL, self).__init__(params, auto_init=False)
78        self.result = []
79
80    def to_screen(self, s, skip_eol=None):
81        print(s)
82
83    def trouble(self, s, tb=None):
84        raise Exception(s)
85
86    def download(self, x):
87        self.result.append(x)
88
89    def expect_warning(self, regex):
90        # Silence an expected warning matching a regex
91        old_report_warning = self.report_warning
92
93        def report_warning(self, message):
94            if re.match(regex, message):
95                return
96            old_report_warning(message)
97        self.report_warning = types.MethodType(report_warning, self)
98
99
100def gettestcases(include_onlymatching=False):
101    for ie in yt_dlp.extractor.gen_extractors():
102        for tc in ie.get_testcases(include_onlymatching):
103            yield tc
104
105
106md5 = lambda s: hashlib.md5(s.encode('utf-8')).hexdigest()
107
108
109def expect_value(self, got, expected, field):
110    if isinstance(expected, compat_str) and expected.startswith('re:'):
111        match_str = expected[len('re:'):]
112        match_rex = re.compile(match_str)
113
114        self.assertTrue(
115            isinstance(got, compat_str),
116            'Expected a %s object, but got %s for field %s' % (
117                compat_str.__name__, type(got).__name__, field))
118        self.assertTrue(
119            match_rex.match(got),
120            'field %s (value: %r) should match %r' % (field, got, match_str))
121    elif isinstance(expected, compat_str) and expected.startswith('startswith:'):
122        start_str = expected[len('startswith:'):]
123        self.assertTrue(
124            isinstance(got, compat_str),
125            'Expected a %s object, but got %s for field %s' % (
126                compat_str.__name__, type(got).__name__, field))
127        self.assertTrue(
128            got.startswith(start_str),
129            'field %s (value: %r) should start with %r' % (field, got, start_str))
130    elif isinstance(expected, compat_str) and expected.startswith('contains:'):
131        contains_str = expected[len('contains:'):]
132        self.assertTrue(
133            isinstance(got, compat_str),
134            'Expected a %s object, but got %s for field %s' % (
135                compat_str.__name__, type(got).__name__, field))
136        self.assertTrue(
137            contains_str in got,
138            'field %s (value: %r) should contain %r' % (field, got, contains_str))
139    elif isinstance(expected, type):
140        self.assertTrue(
141            isinstance(got, expected),
142            'Expected type %r for field %s, but got value %r of type %r' % (expected, field, got, type(got)))
143    elif isinstance(expected, dict) and isinstance(got, dict):
144        expect_dict(self, got, expected)
145    elif isinstance(expected, list) and isinstance(got, list):
146        self.assertEqual(
147            len(expected), len(got),
148            'Expect a list of length %d, but got a list of length %d for field %s' % (
149                len(expected), len(got), field))
150        for index, (item_got, item_expected) in enumerate(zip(got, expected)):
151            type_got = type(item_got)
152            type_expected = type(item_expected)
153            self.assertEqual(
154                type_expected, type_got,
155                'Type mismatch for list item at index %d for field %s, expected %r, got %r' % (
156                    index, field, type_expected, type_got))
157            expect_value(self, item_got, item_expected, field)
158    else:
159        if isinstance(expected, compat_str) and expected.startswith('md5:'):
160            self.assertTrue(
161                isinstance(got, compat_str),
162                'Expected field %s to be a unicode object, but got value %r of type %r' % (field, got, type(got)))
163            got = 'md5:' + md5(got)
164        elif isinstance(expected, compat_str) and re.match(r'^(?:min|max)?count:\d+', expected):
165            self.assertTrue(
166                isinstance(got, (list, dict)),
167                'Expected field %s to be a list or a dict, but it is of type %s' % (
168                    field, type(got).__name__))
169            op, _, expected_num = expected.partition(':')
170            expected_num = int(expected_num)
171            if op == 'mincount':
172                assert_func = assertGreaterEqual
173                msg_tmpl = 'Expected %d items in field %s, but only got %d'
174            elif op == 'maxcount':
175                assert_func = assertLessEqual
176                msg_tmpl = 'Expected maximum %d items in field %s, but got %d'
177            elif op == 'count':
178                assert_func = assertEqual
179                msg_tmpl = 'Expected exactly %d items in field %s, but got %d'
180            else:
181                assert False
182            assert_func(
183                self, len(got), expected_num,
184                msg_tmpl % (expected_num, field, len(got)))
185            return
186        self.assertEqual(
187            expected, got,
188            'Invalid value for field %s, expected %r, got %r' % (field, expected, got))
189
190
191def expect_dict(self, got_dict, expected_dict):
192    for info_field, expected in expected_dict.items():
193        got = got_dict.get(info_field)
194        expect_value(self, got, expected, info_field)
195
196
197def sanitize_got_info_dict(got_dict):
198    IGNORED_FIELDS = (
199        # Format keys
200        'url', 'manifest_url', 'format', 'format_id', 'format_note', 'width', 'height', 'resolution',
201        'dynamic_range', 'tbr', 'abr', 'acodec', 'asr', 'vbr', 'fps', 'vcodec', 'container', 'filesize',
202        'filesize_approx', 'player_url', 'protocol', 'fragment_base_url', 'fragments', 'preference',
203        'language', 'language_preference', 'quality', 'source_preference', 'http_headers',
204        'stretched_ratio', 'no_resume', 'has_drm', 'downloader_options',
205
206        # RTMP formats
207        'page_url', 'app', 'play_path', 'tc_url', 'flash_version', 'rtmp_live', 'rtmp_conn', 'rtmp_protocol', 'rtmp_real_time',
208
209        # Lists
210        'formats', 'thumbnails', 'subtitles', 'automatic_captions', 'comments', 'entries',
211
212        # Auto-generated
213        'autonumber', 'playlist', 'format_index', 'video_ext', 'audio_ext', 'duration_string', 'epoch',
214        'fulltitle', 'extractor', 'extractor_key', 'filepath', 'infojson_filename', 'original_url',
215
216        # Only live_status needs to be checked
217        'is_live', 'was_live',
218    )
219
220    IGNORED_PREFIXES = ('', 'playlist', 'requested', 'webpage')
221
222    def sanitize(key, value):
223        if isinstance(value, str) and len(value) > 100:
224            return f'md5:{md5(value)}'
225        elif isinstance(value, list) and len(value) > 10:
226            return f'count:{len(value)}'
227        return value
228
229    test_info_dict = {
230        key: sanitize(key, value) for key, value in got_dict.items()
231        if value is not None and key not in IGNORED_FIELDS and not any(
232            key.startswith(f'{prefix}_') for prefix in IGNORED_PREFIXES)
233    }
234
235    # display_id may be generated from id
236    if test_info_dict.get('display_id') == test_info_dict['id']:
237        test_info_dict.pop('display_id')
238
239    return test_info_dict
240
241
242def expect_info_dict(self, got_dict, expected_dict):
243    expect_dict(self, got_dict, expected_dict)
244    # Check for the presence of mandatory fields
245    if got_dict.get('_type') not in ('playlist', 'multi_video'):
246        mandatory_fields = ['id', 'title']
247        if expected_dict.get('ext'):
248            mandatory_fields.extend(('url', 'ext'))
249        for key in mandatory_fields:
250            self.assertTrue(got_dict.get(key), 'Missing mandatory field %s' % key)
251    # Check for mandatory fields that are automatically set by YoutubeDL
252    for key in ['webpage_url', 'extractor', 'extractor_key']:
253        self.assertTrue(got_dict.get(key), 'Missing field: %s' % key)
254
255    test_info_dict = sanitize_got_info_dict(got_dict)
256
257    missing_keys = set(test_info_dict.keys()) - set(expected_dict.keys())
258    if missing_keys:
259        def _repr(v):
260            if isinstance(v, compat_str):
261                return "'%s'" % v.replace('\\', '\\\\').replace("'", "\\'").replace('\n', '\\n')
262            else:
263                return repr(v)
264        info_dict_str = ''
265        if len(missing_keys) != len(expected_dict):
266            info_dict_str += ''.join(
267                '    %s: %s,\n' % (_repr(k), _repr(v))
268                for k, v in test_info_dict.items() if k not in missing_keys)
269
270            if info_dict_str:
271                info_dict_str += '\n'
272        info_dict_str += ''.join(
273            '    %s: %s,\n' % (_repr(k), _repr(test_info_dict[k]))
274            for k in missing_keys)
275        write_string(
276            '\n\'info_dict\': {\n' + info_dict_str + '},\n', out=sys.stderr)
277        self.assertFalse(
278            missing_keys,
279            'Missing keys in test definition: %s' % (
280                ', '.join(sorted(missing_keys))))
281
282
283def assertRegexpMatches(self, text, regexp, msg=None):
284    if hasattr(self, 'assertRegexp'):
285        return self.assertRegexp(text, regexp, msg)
286    else:
287        m = re.match(regexp, text)
288        if not m:
289            note = 'Regexp didn\'t match: %r not found' % (regexp)
290            if len(text) < 1000:
291                note += ' in %r' % text
292            if msg is None:
293                msg = note
294            else:
295                msg = note + ', ' + msg
296            self.assertTrue(m, msg)
297
298
299def assertGreaterEqual(self, got, expected, msg=None):
300    if not (got >= expected):
301        if msg is None:
302            msg = '%r not greater than or equal to %r' % (got, expected)
303        self.assertTrue(got >= expected, msg)
304
305
306def assertLessEqual(self, got, expected, msg=None):
307    if not (got <= expected):
308        if msg is None:
309            msg = '%r not less than or equal to %r' % (got, expected)
310        self.assertTrue(got <= expected, msg)
311
312
313def assertEqual(self, got, expected, msg=None):
314    if not (got == expected):
315        if msg is None:
316            msg = '%r not equal to %r' % (got, expected)
317        self.assertTrue(got == expected, msg)
318
319
320def expect_warnings(ydl, warnings_re):
321    real_warning = ydl.report_warning
322
323    def _report_warning(w):
324        if not any(re.search(w_re, w) for w_re in warnings_re):
325            real_warning(w)
326
327    ydl.report_warning = _report_warning
328
329
330def http_server_port(httpd):
331    if os.name == 'java' and isinstance(httpd.socket, ssl.SSLSocket):
332        # In Jython SSLSocket is not a subclass of socket.socket
333        sock = httpd.socket.sock
334    else:
335        sock = httpd.socket
336    return sock.getsockname()[1]
337