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 import urllib.request
  12
  13 from test.helper import FakeYDL, get_params, is_download_test
  14 from yt_dlp.compat import compat_str
  15
  16
  17 @is_download_test
  18 class TestMultipleSocks(unittest.TestCase):
  19     @staticmethod
  20     def _check_params(attrs):
  21         params = get_params()
  22         for attr in attrs:
  23             if attr not in params:
  24                 print('Missing %s. Skipping.' % attr)
  25                 return
  26         return params
  27
  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(),
  37             params['primary_server_ip'])
  38
  39     def test_proxy_https(self):
  40         params = self._check_params(['primary_proxy', 'primary_server_ip'])
  41         if params is None:
  42             return
  43         ydl = FakeYDL({
  44             'proxy': params['primary_proxy']
  45         })
  46         self.assertEqual(
  47             ydl.urlopen('https://yt-dl.org/ip').read().decode(),
  48             params['primary_server_ip'])
  49
  50     def test_secondary_proxy_http(self):
  51         params = self._check_params(['secondary_proxy', 'secondary_server_ip'])
  52         if params is None:
  53             return
  54         ydl = FakeYDL()
  55         req = urllib.request.Request('http://yt-dl.org/ip')
  56         req.add_header('Ytdl-request-proxy', params['secondary_proxy'])
  57         self.assertEqual(
  58             ydl.urlopen(req).read().decode(),
  59             params['secondary_server_ip'])
  60
  61     def test_secondary_proxy_https(self):
  62         params = self._check_params(['secondary_proxy', 'secondary_server_ip'])
  63         if params is None:
  64             return
  65         ydl = FakeYDL()
  66         req = urllib.request.Request('https://yt-dl.org/ip')
  67         req.add_header('Ytdl-request-proxy', params['secondary_proxy'])
  68         self.assertEqual(
  69             ydl.urlopen(req).read().decode(),
  70             params['secondary_server_ip'])
  71
  72
  73 @is_download_test
  74 class TestSocks(unittest.TestCase):
  75     _SKIP_SOCKS_TEST = True
  76
  77     def setUp(self):
  78         if self._SKIP_SOCKS_TEST:
  79             return
  80
  81         self.port = random.randint(20000, 30000)
  82         self.server_process = subprocess.Popen([
  83             'srelay', '-f', '-i', '127.0.0.1:%d' % self.port],
  84             stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
  85
  86     def tearDown(self):
  87         if self._SKIP_SOCKS_TEST:
  88             return
  89
  90         self.server_process.terminate()
  91         self.server_process.communicate()
  92
  93     def _get_ip(self, protocol):
  94         if self._SKIP_SOCKS_TEST:
  95             return '127.0.0.1'
  96
  97         ydl = FakeYDL({
  98             'proxy': '%s://127.0.0.1:%d' % (protocol, self.port),
  99         })
 100         return ydl.urlopen('http://yt-dl.org/ip').read().decode()
 101
 102     def test_socks4(self):
 103         self.assertTrue(isinstance(self._get_ip('socks4'), compat_str))
 104
 105     def test_socks4a(self):
 106         self.assertTrue(isinstance(self._get_ip('socks4a'), compat_str))
 107
 108     def test_socks5(self):
 109         self.assertTrue(isinstance(self._get_ip('socks5'), compat_str))
 110
 111
 112 if __name__ == '__main__':
 113     unittest.main()