test/test_socks.py

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