# -*- coding: utf-8 -*-
"""
proxy.py
~~~~~~~~
⚡⚡⚡ Fast, Lightweight, Pluggable, TLS interception capable proxy server focused on
Network monitoring, controls & Application development, testing, debugging.
:copyright: (c) 2013-present by Abhinav Singh and contributors.
:license: BSD, see LICENSE for more details.
"""
import time
import contextlib
from typing import Any, List, Optional, Generator
import unittest
from ..proxy import Proxy
from ..plugin import CacheResponsesPlugin
from ..common.utils import new_socket_connection
from ..common.constants import DEFAULT_TIMEOUT
[docs]class TestCase(unittest.TestCase):
"""Base TestCase class that automatically setup and tear down proxy.py."""
DEFAULT_PROXY_PY_STARTUP_FLAGS = [
'--hostname', '127.0.0.1',
'--port', '0',
'--num-workers', '1',
'--num-acceptors', '1',
]
PROXY: Optional[Proxy] = None
INPUT_ARGS: Optional[List[str]] = None
[docs] @classmethod
def setUpClass(cls) -> None:
cls.INPUT_ARGS = getattr(cls, 'PROXY_PY_STARTUP_FLAGS') \
if hasattr(cls, 'PROXY_PY_STARTUP_FLAGS') \
else cls.DEFAULT_PROXY_PY_STARTUP_FLAGS
cls.PROXY = Proxy(cls.INPUT_ARGS)
cls.PROXY.flags.plugins[b'HttpProxyBasePlugin'].append(
CacheResponsesPlugin,
)
# pylint: disable=unnecessary-dunder-call
cls.PROXY.__enter__()
assert cls.PROXY.acceptors
cls.wait_for_server(cls.PROXY.flags.port)
[docs] @staticmethod
def wait_for_server(
proxy_port: int,
wait_for_seconds: float = DEFAULT_TIMEOUT,
) -> None:
"""Wait for proxy.py server to come up."""
start_time = time.time()
while True:
try:
new_socket_connection(
('localhost', proxy_port),
).close()
break
except ConnectionRefusedError:
time.sleep(0.1)
if time.time() - start_time > wait_for_seconds:
raise TimeoutError(
'Timed out while waiting for proxy.py to start...',
)
[docs] @classmethod
def tearDownClass(cls) -> None:
assert cls.PROXY
cls.PROXY.__exit__(None, None, None)
cls.PROXY = None
cls.INPUT_ARGS = None
[docs] @contextlib.contextmanager
def vcr(self) -> Generator[None, None, None]:
try:
CacheResponsesPlugin.ENABLED.set()
yield
finally:
CacheResponsesPlugin.ENABLED.clear()
[docs] def run(self, result: Optional[unittest.TestResult] = None) -> Any:
super().run(result)