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