test_socks.py 3.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118
  1. #!/usr/bin/env python
  2. # coding: utf-8
  3. from __future__ import unicode_literals
  4. # Allow direct execution
  5. import os
  6. import sys
  7. import unittest
  8. sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
  9. import random
  10. import subprocess
  11. from test.helper import (
  12. FakeYDL,
  13. get_params,
  14. )
  15. from youtube_dl.compat import (
  16. compat_str,
  17. compat_urllib_request,
  18. )
  19. class TestMultipleSocks(unittest.TestCase):
  20. @staticmethod
  21. def _check_params(attrs):
  22. params = get_params()
  23. for attr in attrs:
  24. if attr not in params:
  25. print('Missing %s. Skipping.' % attr)
  26. return
  27. return params
  28. def test_proxy_http(self):
  29. params = self._check_params(['primary_proxy', 'primary_server_ip'])
  30. if params is None:
  31. return
  32. ydl = FakeYDL({
  33. 'proxy': params['primary_proxy']
  34. })
  35. self.assertEqual(
  36. ydl.urlopen('http://yt-dl.org/ip').read().decode('utf-8'),
  37. params['primary_server_ip'])
  38. def test_proxy_https(self):
  39. params = self._check_params(['primary_proxy', 'primary_server_ip'])
  40. if params is None:
  41. return
  42. ydl = FakeYDL({
  43. 'proxy': params['primary_proxy']
  44. })
  45. self.assertEqual(
  46. ydl.urlopen('https://yt-dl.org/ip').read().decode('utf-8'),
  47. params['primary_server_ip'])
  48. def test_secondary_proxy_http(self):
  49. params = self._check_params(['secondary_proxy', 'secondary_server_ip'])
  50. if params is None:
  51. return
  52. ydl = FakeYDL()
  53. req = compat_urllib_request.Request('http://yt-dl.org/ip')
  54. req.add_header('Ytdl-request-proxy', params['secondary_proxy'])
  55. self.assertEqual(
  56. ydl.urlopen(req).read().decode('utf-8'),
  57. params['secondary_server_ip'])
  58. def test_secondary_proxy_https(self):
  59. params = self._check_params(['secondary_proxy', 'secondary_server_ip'])
  60. if params is None:
  61. return
  62. ydl = FakeYDL()
  63. req = compat_urllib_request.Request('https://yt-dl.org/ip')
  64. req.add_header('Ytdl-request-proxy', params['secondary_proxy'])
  65. self.assertEqual(
  66. ydl.urlopen(req).read().decode('utf-8'),
  67. params['secondary_server_ip'])
  68. class TestSocks(unittest.TestCase):
  69. _SKIP_SOCKS_TEST = True
  70. def setUp(self):
  71. if self._SKIP_SOCKS_TEST:
  72. return
  73. self.port = random.randint(20000, 30000)
  74. self.server_process = subprocess.Popen([
  75. 'srelay', '-f', '-i', '127.0.0.1:%d' % self.port],
  76. stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
  77. def tearDown(self):
  78. if self._SKIP_SOCKS_TEST:
  79. return
  80. self.server_process.terminate()
  81. self.server_process.communicate()
  82. def _get_ip(self, protocol):
  83. if self._SKIP_SOCKS_TEST:
  84. return '127.0.0.1'
  85. ydl = FakeYDL({
  86. 'proxy': '%s://127.0.0.1:%d' % (protocol, self.port),
  87. })
  88. return ydl.urlopen('http://yt-dl.org/ip').read().decode('utf-8')
  89. def test_socks4(self):
  90. self.assertTrue(isinstance(self._get_ip('socks4'), compat_str))
  91. def test_socks4a(self):
  92. self.assertTrue(isinstance(self._get_ip('socks4a'), compat_str))
  93. def test_socks5(self):
  94. self.assertTrue(isinstance(self._get_ip('socks5'), compat_str))
  95. if __name__ == '__main__':
  96. unittest.main()