test_networking_utils.py 7.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208
  1. #!/usr/bin/env python3
  2. # Allow direct execution
  3. import os
  4. import sys
  5. import pytest
  6. sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
  7. import io
  8. import random
  9. import ssl
  10. from yt_dlp.cookies import YoutubeDLCookieJar
  11. from yt_dlp.dependencies import certifi
  12. from yt_dlp.networking import Response
  13. from yt_dlp.networking._helper import (
  14. InstanceStoreMixin,
  15. add_accept_encoding_header,
  16. get_redirect_method,
  17. make_socks_proxy_opts,
  18. select_proxy,
  19. ssl_load_certs,
  20. )
  21. from yt_dlp.networking.exceptions import (
  22. HTTPError,
  23. IncompleteRead,
  24. )
  25. from yt_dlp.socks import ProxyType
  26. from yt_dlp.utils.networking import HTTPHeaderDict
  27. TEST_DIR = os.path.dirname(os.path.abspath(__file__))
  28. class TestNetworkingUtils:
  29. def test_select_proxy(self):
  30. proxies = {
  31. 'all': 'socks5://example.com',
  32. 'http': 'http://example.com:1080',
  33. 'no': 'bypass.example.com,yt-dl.org',
  34. }
  35. assert select_proxy('https://example.com', proxies) == proxies['all']
  36. assert select_proxy('http://example.com', proxies) == proxies['http']
  37. assert select_proxy('http://bypass.example.com', proxies) is None
  38. assert select_proxy('https://yt-dl.org', proxies) is None
  39. @pytest.mark.parametrize('socks_proxy,expected', [
  40. ('socks5h://example.com', {
  41. 'proxytype': ProxyType.SOCKS5,
  42. 'addr': 'example.com',
  43. 'port': 1080,
  44. 'rdns': True,
  45. 'username': None,
  46. 'password': None,
  47. }),
  48. ('socks5://user:@example.com:5555', {
  49. 'proxytype': ProxyType.SOCKS5,
  50. 'addr': 'example.com',
  51. 'port': 5555,
  52. 'rdns': False,
  53. 'username': 'user',
  54. 'password': '',
  55. }),
  56. ('socks4://u%40ser:pa%20ss@127.0.0.1:1080', {
  57. 'proxytype': ProxyType.SOCKS4,
  58. 'addr': '127.0.0.1',
  59. 'port': 1080,
  60. 'rdns': False,
  61. 'username': 'u@ser',
  62. 'password': 'pa ss',
  63. }),
  64. ('socks4a://:pa%20ss@127.0.0.1', {
  65. 'proxytype': ProxyType.SOCKS4A,
  66. 'addr': '127.0.0.1',
  67. 'port': 1080,
  68. 'rdns': True,
  69. 'username': '',
  70. 'password': 'pa ss',
  71. }),
  72. ])
  73. def test_make_socks_proxy_opts(self, socks_proxy, expected):
  74. assert make_socks_proxy_opts(socks_proxy) == expected
  75. def test_make_socks_proxy_unknown(self):
  76. with pytest.raises(ValueError, match='Unknown SOCKS proxy version: socks'):
  77. make_socks_proxy_opts('socks://127.0.0.1')
  78. @pytest.mark.skipif(not certifi, reason='certifi is not installed')
  79. def test_load_certifi(self):
  80. context_certifi = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
  81. context_certifi.load_verify_locations(cafile=certifi.where())
  82. context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
  83. ssl_load_certs(context, use_certifi=True)
  84. assert context.get_ca_certs() == context_certifi.get_ca_certs()
  85. context_default = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
  86. context_default.load_default_certs()
  87. context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
  88. ssl_load_certs(context, use_certifi=False)
  89. assert context.get_ca_certs() == context_default.get_ca_certs()
  90. if context_default.get_ca_certs() == context_certifi.get_ca_certs():
  91. pytest.skip('System uses certifi as default. The test is not valid')
  92. @pytest.mark.parametrize('method,status,expected', [
  93. ('GET', 303, 'GET'),
  94. ('HEAD', 303, 'HEAD'),
  95. ('PUT', 303, 'GET'),
  96. ('POST', 301, 'GET'),
  97. ('HEAD', 301, 'HEAD'),
  98. ('POST', 302, 'GET'),
  99. ('HEAD', 302, 'HEAD'),
  100. ('PUT', 302, 'PUT'),
  101. ('POST', 308, 'POST'),
  102. ('POST', 307, 'POST'),
  103. ('HEAD', 308, 'HEAD'),
  104. ('HEAD', 307, 'HEAD'),
  105. ])
  106. def test_get_redirect_method(self, method, status, expected):
  107. assert get_redirect_method(method, status) == expected
  108. @pytest.mark.parametrize('headers,supported_encodings,expected', [
  109. ({'Accept-Encoding': 'br'}, ['gzip', 'br'], {'Accept-Encoding': 'br'}),
  110. ({}, ['gzip', 'br'], {'Accept-Encoding': 'gzip, br'}),
  111. ({'Content-type': 'application/json'}, [], {'Content-type': 'application/json', 'Accept-Encoding': 'identity'}),
  112. ])
  113. def test_add_accept_encoding_header(self, headers, supported_encodings, expected):
  114. headers = HTTPHeaderDict(headers)
  115. add_accept_encoding_header(headers, supported_encodings)
  116. assert headers == HTTPHeaderDict(expected)
  117. class TestInstanceStoreMixin:
  118. class FakeInstanceStoreMixin(InstanceStoreMixin):
  119. def _create_instance(self, **kwargs):
  120. return random.randint(0, 1000000)
  121. def _close_instance(self, instance):
  122. pass
  123. def test_mixin(self):
  124. mixin = self.FakeInstanceStoreMixin()
  125. assert mixin._get_instance(d={'a': 1, 'b': 2, 'c': {'d', 4}}) == mixin._get_instance(d={'a': 1, 'b': 2, 'c': {'d', 4}})
  126. assert mixin._get_instance(d={'a': 1, 'b': 2, 'c': {'e', 4}}) != mixin._get_instance(d={'a': 1, 'b': 2, 'c': {'d', 4}})
  127. assert mixin._get_instance(d={'a': 1, 'b': 2, 'c': {'d', 4}} != mixin._get_instance(d={'a': 1, 'b': 2, 'g': {'d', 4}}))
  128. assert mixin._get_instance(d={'a': 1}, e=[1, 2, 3]) == mixin._get_instance(d={'a': 1}, e=[1, 2, 3])
  129. assert mixin._get_instance(d={'a': 1}, e=[1, 2, 3]) != mixin._get_instance(d={'a': 1}, e=[1, 2, 3, 4])
  130. cookiejar = YoutubeDLCookieJar()
  131. assert mixin._get_instance(b=[1, 2], c=cookiejar) == mixin._get_instance(b=[1, 2], c=cookiejar)
  132. assert mixin._get_instance(b=[1, 2], c=cookiejar) != mixin._get_instance(b=[1, 2], c=YoutubeDLCookieJar())
  133. # Different order
  134. assert mixin._get_instance(c=cookiejar, b=[1, 2]) == mixin._get_instance(b=[1, 2], c=cookiejar)
  135. m = mixin._get_instance(t=1234)
  136. assert mixin._get_instance(t=1234) == m
  137. mixin._clear_instances()
  138. assert mixin._get_instance(t=1234) != m
  139. class TestNetworkingExceptions:
  140. @staticmethod
  141. def create_response(status):
  142. return Response(fp=io.BytesIO(b'test'), url='http://example.com', headers={'tesT': 'test'}, status=status)
  143. def test_http_error(self):
  144. response = self.create_response(403)
  145. error = HTTPError(response)
  146. assert error.status == 403
  147. assert str(error) == error.msg == 'HTTP Error 403: Forbidden'
  148. assert error.reason == response.reason
  149. assert error.response is response
  150. data = error.response.read()
  151. assert data == b'test'
  152. assert repr(error) == '<HTTPError 403: Forbidden>'
  153. def test_redirect_http_error(self):
  154. response = self.create_response(301)
  155. error = HTTPError(response, redirect_loop=True)
  156. assert str(error) == error.msg == 'HTTP Error 301: Moved Permanently (redirect loop detected)'
  157. assert error.reason == 'Moved Permanently'
  158. def test_incomplete_read_error(self):
  159. error = IncompleteRead(4, 3, cause='test')
  160. assert isinstance(error, IncompleteRead)
  161. assert repr(error) == '<IncompleteRead: 4 bytes read, 3 more expected>'
  162. assert str(error) == error.msg == '4 bytes read, 3 more expected'
  163. assert error.partial == 4
  164. assert error.expected == 3
  165. assert error.cause == 'test'
  166. error = IncompleteRead(3)
  167. assert repr(error) == '<IncompleteRead: 3 bytes read>'
  168. assert str(error) == '3 bytes read'