diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 0000000..fb67425 --- /dev/null +++ b/pytest.ini @@ -0,0 +1,4 @@ +# pytest.ini +[pytest] +testpaths = + tests diff --git a/requirements-dev.txt b/requirements-dev.txt new file mode 100644 index 0000000..559a4fc --- /dev/null +++ b/requirements-dev.txt @@ -0,0 +1 @@ +pytest>=6.2.1 diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..2694317 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,14 @@ +import pytest +import urllib3 +import urllib +import urllib.request +import socket + +# https://realpython.com/pytest-python-testing/ +@pytest.fixture(autouse=True) +def disable_network_calls(monkeypatch): + def stunted_get(*args, **kwargs): + raise RuntimeError('Network access not allowed during testing!') + monkeypatch.setattr(urllib.request, 'Request', stunted_get) + monkeypatch.setattr(urllib3.PoolManager, 'request', stunted_get) + monkeypatch.setattr(socket, 'socket', stunted_get) diff --git a/tests/test_responses/429.html b/tests/test_responses/429.html new file mode 100644 index 0000000..9bde0f9 --- /dev/null +++ b/tests/test_responses/429.html @@ -0,0 +1,28 @@ + + +https://m.youtube.com/watch?v=aaaaaaaaaaa&pbj=1&bpctr=9999999999 + +
+

+
+ + +
+ +
+
+ +
+About this page

+ +Our systems have detected unusual traffic from your computer network. This page checks to see if it's really you sending the requests, and not a robot. Why did this happen?

+ + + +IP address: 2001:4ba0:ffea:2ae::a10
Time: 2019-12-21T04:28:41Z
URL: https://m.youtube.com/watch?v=aaaaaaaaaaa&pbj=1&bpctr=9999999999
+
+
+ + diff --git a/tests/test_util.py b/tests/test_util.py new file mode 100644 index 0000000..bc10de4 --- /dev/null +++ b/tests/test_util.py @@ -0,0 +1,76 @@ +from youtube import util +import settings +import pytest # overview: https://realpython.com/pytest-python-testing/ +import urllib3 +import io +import os +import stem + + +def load_test_page(name): + with open(os.path.join('./tests/test_responses', name), 'rb') as f: + return f.read() + + +html429 = load_test_page('429.html') + + +class MockResponse(urllib3.response.HTTPResponse): + def __init__(self, body='success', headers=None, status=200, reason=''): + print(body[0:10]) + headers = headers or {} + if isinstance(body, str): + body = body.encode('utf-8') + self.body_io = io.BytesIO(body) + self.read = self.body_io.read + urllib3.response.HTTPResponse.__init__( + self, body=body, headers=headers, status=status, + preload_content=False, decode_content=False, reason=reason + ) + + +class NewIdentityState(): + MAX_TRIES = util.TorManager.MAX_TRIES + def __init__(self, new_identities_till_success): + self.new_identities_till_success = new_identities_till_success + + def new_identity(self, *args, **kwargs): + print('newidentity') + self.new_identities_till_success -= 1 + + def fetch_url_response(self, *args, **kwargs): + cleanup_func = (lambda r: None) + if self.new_identities_till_success == 0: + return MockResponse(), cleanup_func + return MockResponse(body=html429, status=429), cleanup_func + + +class MockController(): + def authenticate(self, *args, **kwargs): + pass + @classmethod + def from_port(cls, *args, **kwargs): + return cls() + def __enter__(self, *args, **kwargs): + return self + def __exit__(self, *args, **kwargs): + pass + + +@pytest.mark.parametrize('new_identities_till_success', + [i for i in range(0, NewIdentityState.MAX_TRIES+2)]) +def test_exit_node_retry(monkeypatch, new_identities_till_success): + new_identity_state = NewIdentityState(new_identities_till_success) + # https://docs.pytest.org/en/stable/monkeypatch.html + monkeypatch.setattr(settings, 'route_tor', 1) + monkeypatch.setattr(util, 'tor_manager', util.TorManager()) # fresh one + MockController.signal = new_identity_state.new_identity + monkeypatch.setattr(stem.control, 'Controller', MockController) + monkeypatch.setattr(util, 'fetch_url_response', + new_identity_state.fetch_url_response) + if new_identities_till_success <= NewIdentityState.MAX_TRIES: + assert util.fetch_url('url') == b'success' + else: + with pytest.raises(util.FetchError) as excinfo: + util.fetch_url('url') + assert int(excinfo.value.code) == 429 diff --git a/youtube/util.py b/youtube/util.py index 355d8c7..1544c94 100644 --- a/youtube/util.py +++ b/youtube/util.py @@ -60,6 +60,9 @@ connection_pool = urllib3.PoolManager(cert_reqs='CERT_REQUIRED') class TorManager: + MAX_TRIES = 3 + COOLDOWN_TIME = 5 + def __init__(self): self.old_tor_connection_pool = None self.tor_connection_pool = urllib3.contrib.socks.SOCKSProxyManager( @@ -69,6 +72,7 @@ class TorManager: self.new_identity_lock = gevent.lock.BoundedSemaphore(1) self.last_new_identity_time = time.monotonic() - 20 + self.try_num = 1 def refresh_tor_connection_pool(self): self.tor_connection_pool.clear() @@ -108,9 +112,14 @@ class TorManager: return None delta = time.monotonic() - self.last_new_identity_time - if delta < 20: - print('new_identity: Retried already within last 20 seconds') - return 'Retried with new circuit once (max) within last 20 seconds.' + if delta < self.COOLDOWN_TIME and self.try_num == 1: + err = ('Retried with new circuit %d times (max) within last ' + '%d seconds.' % (self.MAX_TRIES, self.COOLDOWN_TIME)) + print('new_identity:', err) + return err + elif delta >= self.COOLDOWN_TIME: + self.try_num = 1 + try: port = settings.tor_control_port with stem.control.Controller.from_port(port=port) as controller: @@ -120,10 +129,14 @@ class TorManager: print('new_identity: NEWNYM signal sent') self.last_new_identity_time = time.monotonic() self.refresh_tor_connection_pool() - return None except stem.SocketError: traceback.print_exc() return 'Failed to connect to Tor control port.' + finally: + self.try_num += 1 + if self.try_num > self.MAX_TRIES: + self.try_num = 1 + return None finally: self.new_identity_lock.release()