test/test_socks.py

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