test/test_networking_utils.py

   1 #!/usr/bin/env python3
   2
   3 # Allow direct execution
   4 import os
   5 import sys
   6
   7 import pytest
   8
   9 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
  10
  11 import io
  12 import platform
  13 import random
  14 import ssl
  15 import urllib.error
  16
  17 from yt_dlp.cookies import YoutubeDLCookieJar
  18 from yt_dlp.dependencies import certifi
  19 from yt_dlp.networking import Response
  20 from yt_dlp.networking._helper import (
  21     InstanceStoreMixin,
  22     add_accept_encoding_header,
  23     get_redirect_method,
  24     make_socks_proxy_opts,
  25     select_proxy,
  26     ssl_load_certs,
  27 )
  28 from yt_dlp.networking.exceptions import (
  29     HTTPError,
  30     IncompleteRead,
  31     _CompatHTTPError,
  32 )
  33 from yt_dlp.socks import ProxyType
  34 from yt_dlp.utils.networking import HTTPHeaderDict
  35
  36 TEST_DIR = os.path.dirname(os.path.abspath(__file__))
  37
  38
  39 class TestNetworkingUtils:
  40
  41     def test_select_proxy(self):
  42         proxies = {
  43             'all': 'socks5://example.com',
  44             'http': 'http://example.com:1080',
  45             'no': 'bypass.example.com,yt-dl.org'
  46         }
  47
  48         assert select_proxy('https://example.com', proxies) == proxies['all']
  49         assert select_proxy('http://example.com', proxies) == proxies['http']
  50         assert select_proxy('http://bypass.example.com', proxies) is None
  51         assert select_proxy('https://yt-dl.org', proxies) is None
  52
  53     @pytest.mark.parametrize('socks_proxy,expected', [
  54         ('socks5h://example.com', {
  55             'proxytype': ProxyType.SOCKS5,
  56             'addr': 'example.com',
  57             'port': 1080,
  58             'rdns': True,
  59             'username': None,
  60             'password': None
  61         }),
  62         ('socks5://user:@example.com:5555', {
  63             'proxytype': ProxyType.SOCKS5,
  64             'addr': 'example.com',
  65             'port': 5555,
  66             'rdns': False,
  67             'username': 'user',
  68             'password': ''
  69         }),
  70         ('socks4://u%40ser:pa%20ss@127.0.0.1:1080', {
  71             'proxytype': ProxyType.SOCKS4,
  72             'addr': '127.0.0.1',
  73             'port': 1080,
  74             'rdns': False,
  75             'username': 'u@ser',
  76             'password': 'pa ss'
  77         }),
  78         ('socks4a://:pa%20ss@127.0.0.1', {
  79             'proxytype': ProxyType.SOCKS4A,
  80             'addr': '127.0.0.1',
  81             'port': 1080,
  82             'rdns': True,
  83             'username': '',
  84             'password': 'pa ss'
  85         })
  86     ])
  87     def test_make_socks_proxy_opts(self, socks_proxy, expected):
  88         assert make_socks_proxy_opts(socks_proxy) == expected
  89
  90     def test_make_socks_proxy_unknown(self):
  91         with pytest.raises(ValueError, match='Unknown SOCKS proxy version: socks'):
  92             make_socks_proxy_opts('socks://127.0.0.1')
  93
  94     @pytest.mark.skipif(not certifi, reason='certifi is not installed')
  95     def test_load_certifi(self):
  96         context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
  97         context2 = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
  98         ssl_load_certs(context, use_certifi=True)
  99         context2.load_verify_locations(cafile=certifi.where())
 100         assert context.get_ca_certs() == context2.get_ca_certs()
 101
 102         # Test load normal certs
 103         # XXX: could there be a case where system certs are the same as certifi?
 104         context3 = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
 105         ssl_load_certs(context3, use_certifi=False)
 106         assert context3.get_ca_certs() != context.get_ca_certs()
 107
 108     @pytest.mark.parametrize('method,status,expected', [
 109         ('GET', 303, 'GET'),
 110         ('HEAD', 303, 'HEAD'),
 111         ('PUT', 303, 'GET'),
 112         ('POST', 301, 'GET'),
 113         ('HEAD', 301, 'HEAD'),
 114         ('POST', 302, 'GET'),
 115         ('HEAD', 302, 'HEAD'),
 116         ('PUT', 302, 'PUT'),
 117         ('POST', 308, 'POST'),
 118         ('POST', 307, 'POST'),
 119         ('HEAD', 308, 'HEAD'),
 120         ('HEAD', 307, 'HEAD'),
 121     ])
 122     def test_get_redirect_method(self, method, status, expected):
 123         assert get_redirect_method(method, status) == expected
 124
 125     @pytest.mark.parametrize('headers,supported_encodings,expected', [
 126         ({'Accept-Encoding': 'br'}, ['gzip', 'br'], {'Accept-Encoding': 'br'}),
 127         ({}, ['gzip', 'br'], {'Accept-Encoding': 'gzip, br'}),
 128         ({'Content-type': 'application/json'}, [], {'Content-type': 'application/json', 'Accept-Encoding': 'identity'}),
 129     ])
 130     def test_add_accept_encoding_header(self, headers, supported_encodings, expected):
 131         headers = HTTPHeaderDict(headers)
 132         add_accept_encoding_header(headers, supported_encodings)
 133         assert headers == HTTPHeaderDict(expected)
 134
 135
 136 class TestInstanceStoreMixin:
 137
 138     class FakeInstanceStoreMixin(InstanceStoreMixin):
 139         def _create_instance(self, **kwargs):
 140             return random.randint(0, 1000000)
 141
 142         def _close_instance(self, instance):
 143             pass
 144
 145     def test_mixin(self):
 146         mixin = self.FakeInstanceStoreMixin()
 147         assert mixin._get_instance(d={'a': 1, 'b': 2, 'c': {'d', 4}}) == mixin._get_instance(d={'a': 1, 'b': 2, 'c': {'d', 4}})
 148
 149         assert mixin._get_instance(d={'a': 1, 'b': 2, 'c': {'e', 4}}) != mixin._get_instance(d={'a': 1, 'b': 2, 'c': {'d', 4}})
 150
 151         assert mixin._get_instance(d={'a': 1, 'b': 2, 'c': {'d', 4}} != mixin._get_instance(d={'a': 1, 'b': 2, 'g': {'d', 4}}))
 152
 153         assert mixin._get_instance(d={'a': 1}, e=[1, 2, 3]) == mixin._get_instance(d={'a': 1}, e=[1, 2, 3])
 154
 155         assert mixin._get_instance(d={'a': 1}, e=[1, 2, 3]) != mixin._get_instance(d={'a': 1}, e=[1, 2, 3, 4])
 156
 157         cookiejar = YoutubeDLCookieJar()
 158         assert mixin._get_instance(b=[1, 2], c=cookiejar) == mixin._get_instance(b=[1, 2], c=cookiejar)
 159
 160         assert mixin._get_instance(b=[1, 2], c=cookiejar) != mixin._get_instance(b=[1, 2], c=YoutubeDLCookieJar())
 161
 162         # Different order
 163         assert mixin._get_instance(c=cookiejar, b=[1, 2]) == mixin._get_instance(b=[1, 2], c=cookiejar)
 164
 165         m = mixin._get_instance(t=1234)
 166         assert mixin._get_instance(t=1234) == m
 167         mixin._clear_instances()
 168         assert mixin._get_instance(t=1234) != m
 169
 170
 171 class TestNetworkingExceptions:
 172
 173     @staticmethod
 174     def create_response(status):
 175         return Response(fp=io.BytesIO(b'test'), url='http://example.com', headers={'tesT': 'test'}, status=status)
 176
 177     @pytest.mark.parametrize('http_error_class', [HTTPError, lambda r: _CompatHTTPError(HTTPError(r))])
 178     def test_http_error(self, http_error_class):
 179
 180         response = self.create_response(403)
 181         error = http_error_class(response)
 182
 183         assert error.status == 403
 184         assert str(error) == error.msg == 'HTTP Error 403: Forbidden'
 185         assert error.reason == response.reason
 186         assert error.response is response
 187
 188         data = error.response.read()
 189         assert data == b'test'
 190         assert repr(error) == '<HTTPError 403: Forbidden>'
 191
 192     @pytest.mark.parametrize('http_error_class', [HTTPError, lambda *args, **kwargs: _CompatHTTPError(HTTPError(*args, **kwargs))])
 193     def test_redirect_http_error(self, http_error_class):
 194         response = self.create_response(301)
 195         error = http_error_class(response, redirect_loop=True)
 196         assert str(error) == error.msg == 'HTTP Error 301: Moved Permanently (redirect loop detected)'
 197         assert error.reason == 'Moved Permanently'
 198
 199     def test_compat_http_error(self):
 200         response = self.create_response(403)
 201         error = _CompatHTTPError(HTTPError(response))
 202         assert isinstance(error, HTTPError)
 203         assert isinstance(error, urllib.error.HTTPError)
 204
 205         assert error.code == 403
 206         assert error.getcode() == 403
 207         assert error.hdrs is error.response.headers
 208         assert error.info() is error.response.headers
 209         assert error.headers is error.response.headers
 210         assert error.filename == error.response.url
 211         assert error.url == error.response.url
 212         assert error.geturl() == error.response.url
 213
 214         # Passthrough file operations
 215         assert error.read() == b'test'
 216         assert not error.closed
 217         # Technically Response operations are also passed through, which should not be used.
 218         assert error.get_header('test') == 'test'
 219
 220     @pytest.mark.skipif(
 221         platform.python_implementation() == 'PyPy', reason='garbage collector works differently in pypy')
 222     def test_compat_http_error_autoclose(self):
 223         # Compat HTTPError should not autoclose response
 224         response = self.create_response(403)
 225         _CompatHTTPError(HTTPError(response))
 226         assert not response.closed
 227
 228     def test_incomplete_read_error(self):
 229         error = IncompleteRead(b'test', 3, cause='test')
 230         assert isinstance(error, IncompleteRead)
 231         assert repr(error) == '<IncompleteRead: 4 bytes read, 3 more expected>'
 232         assert str(error) == error.msg == '4 bytes read, 3 more expected'
 233         assert error.partial == b'test'
 234         assert error.expected == 3
 235         assert error.cause == 'test'
 236
 237         error = IncompleteRead(b'aaa')
 238         assert repr(error) == '<IncompleteRead: 3 bytes read>'
 239         assert str(error) == '3 bytes read'