Run mitmproxy only once per test case.

This reduces the run time of the test suite by 40%, from 6.7s to 4.1s.

Author: aaugustin

tests/asyncio/test_client.py

@@ -22,7 +22,7 @@
 )
 from websockets.extensions.permessage_deflate import PerMessageDeflate
 
-from ..proxy import async_proxy
+from ..proxy import ProxyMixin
 from ..utils import (
     CLIENT_CONTEXT,
     MS,
@@ -388,7 +388,7 @@ def remove_accept_header(self, request, response):
 
     async def test_timeout_during_handshake(self):
         """Client times out before receiving handshake response from server."""
-        # Replace the WebSocket server with a TCP server that does't respond.
+        # Replace the WebSocket server with a TCP server that doesn't respond.
         with socket.create_server(("localhost", 0)) as sock:
             host, port = sock.getsockname()
             with self.assertRaises(TimeoutError) as raised:
@@ -508,7 +508,7 @@ async def test_reject_invalid_server_certificate(self):
         """Client rejects certificate where server certificate isn't trusted."""
         async with serve(*args, ssl=SERVER_CONTEXT) as server:
             with self.assertRaises(ssl.SSLCertVerificationError) as raised:
-                # The test certificate isn't trusted system-wide.
+                # The test certificate is self-signed.
                 async with connect(get_uri(server)):
                     self.fail("did not raise")
             self.assertIn(
@@ -566,126 +566,104 @@ def redirect(connection, request):
 
 
 @unittest.skipUnless("mitmproxy" in sys.modules, "mitmproxy not installed")
-class ProxyClientTests(unittest.IsolatedAsyncioTestCase):
-    @contextlib.asynccontextmanager
-    async def socks_proxy(self, auth=None):
-        if auth:
-            proxyauth = "hello:iloveyou"
-            proxy_uri = "http://hello:iloveyou@localhost:51080"
-        else:
-            proxyauth = None
-            proxy_uri = "http://localhost:51080"
-        async with async_proxy(
-            mode=["socks5@51080"],
-            proxyauth=proxyauth,
-        ) as record_flows:
-            with patch_environ({"socks_proxy": proxy_uri}):
-                yield record_flows
+class SocksProxyClientTests(ProxyMixin, unittest.IsolatedAsyncioTestCase):
+    proxy_mode = "socks5@51080"
 
     async def test_socks_proxy(self):
         """Client connects to server through a SOCKS5 proxy."""
-        async with self.socks_proxy() as proxy:
+        with patch_environ({"socks_proxy": "http://localhost:51080"}):
             async with serve(*args) as server:
                 async with connect(get_uri(server)) as client:
                     self.assertEqual(client.protocol.state.name, "OPEN")
-        self.assertEqual(len(proxy.get_flows()), 1)
+        self.assertNumFlows(1)
 
     async def test_secure_socks_proxy(self):
         """Client connects to server securely through a SOCKS5 proxy."""
-        async with self.socks_proxy() as proxy:
+        with patch_environ({"socks_proxy": "http://localhost:51080"}):
             async with serve(*args, ssl=SERVER_CONTEXT) as server:
                 async with connect(get_uri(server), ssl=CLIENT_CONTEXT) as client:
                     self.assertEqual(client.protocol.state.name, "OPEN")
-        self.assertEqual(len(proxy.get_flows()), 1)
+        self.assertNumFlows(1)
 
     async def test_authenticated_socks_proxy(self):
         """Client connects to server through an authenticated SOCKS5 proxy."""
-        async with self.socks_proxy(auth=True) as proxy:
-            async with serve(*args) as server:
-                async with connect(get_uri(server)) as client:
-                    self.assertEqual(client.protocol.state.name, "OPEN")
-        self.assertEqual(len(proxy.get_flows()), 1)
+        try:
+            self.proxy_options.update(proxyauth="hello:iloveyou")
+            with patch_environ(
+                {"socks_proxy": "http://hello:iloveyou@localhost:51080"}
+            ):
+                async with serve(*args) as server:
+                    async with connect(get_uri(server)) as client:
+                        self.assertEqual(client.protocol.state.name, "OPEN")
+        finally:
+            self.proxy_options.update(proxyauth=None)
+        self.assertNumFlows(1)
 
-    async def test_socks_proxy_connection_error(self):
-        """Client receives an error when connecting to the SOCKS5 proxy."""
+    async def test_authenticated_socks_proxy_error(self):
+        """Client fails to authenticate to the SOCKS5 proxy."""
         from python_socks import ProxyError as SocksProxyError
 
-        async with self.socks_proxy(auth=True) as proxy:
-            with self.assertRaises(ProxyError) as raised:
-                async with connect(
-                    "ws://example.com/",
-                    proxy="socks5h://localhost:51080",  # remove credentials
-                ):
-                    self.fail("did not raise")
+        try:
+            self.proxy_options.update(proxyauth="any")
+            with patch_environ({"socks_proxy": "http://localhost:51080"}):
+                with self.assertRaises(ProxyError) as raised:
+                    async with connect("ws://example.com/"):
+                        self.fail("did not raise")
+        finally:
+            self.proxy_options.update(proxyauth=None)
         self.assertEqual(
             str(raised.exception),
             "failed to connect to SOCKS proxy",
         )
         self.assertIsInstance(raised.exception.__cause__, SocksProxyError)
-        self.assertEqual(len(proxy.get_flows()), 0)
+        self.assertNumFlows(0)
 
-    async def test_socks_proxy_connection_fails(self):
+    async def test_socks_proxy_connection_failure(self):
         """Client fails to connect to the SOCKS5 proxy."""
         from python_socks import ProxyConnectionError as SocksProxyConnectionError
 
-        with self.assertRaises(OSError) as raised:
-            async with connect(
-                "ws://example.com/",
-                proxy="socks5h://localhost:51080",  # nothing at this address
-            ):
-                self.fail("did not raise")
+        with patch_environ({"socks_proxy": "http://localhost:61080"}):  # bad port
+            with self.assertRaises(OSError) as raised:
+                async with connect("ws://example.com/"):
+                    self.fail("did not raise")
         # Don't test str(raised.exception) because we don't control it.
         self.assertIsInstance(raised.exception, SocksProxyConnectionError)
+        self.assertNumFlows(0)
 
     async def test_socks_proxy_connection_timeout(self):
         """Client times out while connecting to the SOCKS5 proxy."""
-        # Replace the proxy with a TCP server that does't respond.
+        # Replace the proxy with a TCP server that doesn't respond.
         with socket.create_server(("localhost", 0)) as sock:
             host, port = sock.getsockname()
-            with self.assertRaises(TimeoutError) as raised:
-                async with connect(
-                    "ws://example.com/",
-                    proxy=f"socks5h://{host}:{port}/",
-                    open_timeout=MS,
-                ):
-                    self.fail("did not raise")
+            with patch_environ({"socks_proxy": f"http://{host}:{port}"}):
+                with self.assertRaises(TimeoutError) as raised:
+                    async with connect("ws://example.com/", open_timeout=MS):
+                        self.fail("did not raise")
         self.assertEqual(
             str(raised.exception),
             "timed out during handshake",
         )
+        self.assertNumFlows(0)
 
-    async def test_explicit_proxy(self):
-        """Client connects to server through a proxy set explicitly."""
-        async with async_proxy(mode=["socks5@51080"]) as proxy:
-            async with serve(*args) as server:
-                async with connect(
-                    get_uri(server),
-                    # Take this opportunity to test socks5 instead of socks5h.
-                    proxy="socks5://localhost:51080",
-                ) as client:
-                    self.assertEqual(client.protocol.state.name, "OPEN")
-        self.assertEqual(len(proxy.get_flows()), 1)
+    async def test_explicit_socks_proxy(self):
+        """Client connects to server through a SOCKS5 proxy set explicitly."""
+        async with serve(*args) as server:
+            async with connect(
+                get_uri(server),
+                # Take this opportunity to test socks5 instead of socks5h.
+                proxy="socks5://localhost:51080",
+            ) as client:
+                self.assertEqual(client.protocol.state.name, "OPEN")
+        self.assertNumFlows(1)
 
     async def test_ignore_proxy_with_existing_socket(self):
         """Client connects using a pre-existing socket."""
-        async with self.socks_proxy() as proxy:
+        with patch_environ({"socks_proxy": "http://localhost:51080"}):
             async with serve(*args) as server:
                 with socket.create_connection(get_host_port(server)) as sock:
                     # Use a non-existing domain to ensure we connect to sock.
                     async with connect("ws://invalid/", sock=sock) as client:
                         self.assertEqual(client.protocol.state.name, "OPEN")
-        self.assertEqual(len(proxy.get_flows()), 0)
-
-    async def test_unsupported_proxy(self):
-        """Client connects to server through an unsupported proxy."""
-        with patch_environ({"ws_proxy": "other://localhost:51080"}):
-            with self.assertRaises(InvalidProxy) as raised:
-                async with connect("ws://example.com/"):
-                    self.fail("did not raise")
-        self.assertEqual(
-            str(raised.exception),
-            "other://localhost:51080 isn't a valid proxy: scheme other isn't supported",
-        )
 
 
 @unittest.skipUnless(hasattr(socket, "AF_UNIX"), "this test requires Unix sockets")
@@ -724,10 +702,7 @@ def redirect(connection, request):
             "cannot follow cross-origin redirect to ws://other/ with a Unix socket",
         )
 
-
-@unittest.skipUnless(hasattr(socket, "AF_UNIX"), "this test requires Unix sockets")
-class SecureUnixClientTests(unittest.IsolatedAsyncioTestCase):
-    async def test_connection(self):
+    async def test_secure_connection(self):
         """Client connects to server securely over a Unix socket."""
         with temp_unix_socket_path() as path:
             async with unix_serve(handler, path, ssl=SERVER_CONTEXT):
@@ -769,6 +744,16 @@ async def test_secure_uri_without_ssl(self):
             "ssl=None is incompatible with a wss:// URI",
         )
 
+    async def test_unsupported_proxy(self):
+        """Client rejects unsupported proxy."""
+        with self.assertRaises(InvalidProxy) as raised:
+            async with connect("ws://example.com/", proxy="other://localhost:51080"):
+                self.fail("did not raise")
+        self.assertEqual(
+            str(raised.exception),
+            "other://localhost:51080 isn't a valid proxy: scheme other isn't supported",
+        )
+
     async def test_unix_without_path_or_sock(self):
         """Unix client requires path when sock isn't provided."""
         with self.assertRaises(ValueError) as raised:

tests/proxy.py

@@ -1,15 +1,14 @@
 import asyncio
-import contextlib
 import pathlib
 import threading
 import warnings
 
 
-warnings.filterwarnings("ignore", category=DeprecationWarning, module="mitmproxy")
-warnings.filterwarnings("ignore", category=DeprecationWarning, module="passlib")
-warnings.filterwarnings("ignore", category=DeprecationWarning, module="pyasn1")
-
 try:
+    # Ignore deprecation warnings raised by mitmproxy dependencies at import time.
+    warnings.filterwarnings("ignore", category=DeprecationWarning, module="passlib")
+    warnings.filterwarnings("ignore", category=DeprecationWarning, module="pyasn1")
+
     from mitmproxy.addons import core, next_layer, proxyauth, proxyserver, tlsconfig
     from mitmproxy.master import Master
     from mitmproxy.options import Options
@@ -18,72 +17,87 @@
 
 
 class RecordFlows:
-    def __init__(self):
-        self.ready = asyncio.get_running_loop().create_future()
+    def __init__(self, on_running):
+        self.running = on_running
         self.flows = []
 
-    def running(self):
-        self.ready.set_result(None)
-
     def websocket_start(self, flow):
         self.flows.append(flow)
 
     def get_flows(self):
         flows, self.flows[:] = self.flows[:], []
         return flows
 
+    def reset_flows(self):
+        self.flows = []
+
+
+class ProxyMixin:
+    """
+    Run mitmproxy in a background thread.
+
+    While it's uncommon to run two event loops in two threads, tests for the
+    asyncio implementation rely on this class too because it starts an event
+    loop for mitm proxy once, then a new event loop for each test.
+    """
+
+    proxy_mode = None
+
+    @classmethod
+    async def run_proxy(cls):
+        cls.proxy_loop = loop = asyncio.get_event_loop()
+        cls.proxy_stop = stop = loop.create_future()
+
+        cls.proxy_options = options = Options(mode=[cls.proxy_mode])
+        cls.proxy_master = master = Master(options)
+        master.addons.add(
+            core.Core(),
+            proxyauth.ProxyAuth(),
+            proxyserver.Proxyserver(),
+            next_layer.NextLayer(),
+            tlsconfig.TlsConfig(),
+            RecordFlows(on_running=cls.proxy_ready.set),
+        )
+        options.update(
+            # Use test certificate for TLS between client and proxy.
+            certs=[str(pathlib.Path(__file__).with_name("test_localhost.pem"))],
+            # Disable TLS verification between proxy and upstream.
+            ssl_insecure=True,
+        )
+
+        task = loop.create_task(cls.proxy_master.run())
+        await stop
 
-@contextlib.asynccontextmanager
-async def async_proxy(mode, **config):
-    options = Options(mode=mode)
-    master = Master(options)
-    record_flows = RecordFlows()
-    master.addons.add(
-        core.Core(),
-        proxyauth.ProxyAuth(),
-        proxyserver.Proxyserver(),
-        next_layer.NextLayer(),
-        tlsconfig.TlsConfig(),
-        record_flows,
-    )
-    config.update(
-        # Use our test certificate for TLS between client and proxy
-        # and disable TLS verification between proxy and upstream.
-        certs=[str(pathlib.Path(__file__).with_name("test_localhost.pem"))],
-        ssl_insecure=True,
-    )
-    options.update(**config)
-
-    asyncio.create_task(master.run())
-    try:
-        await record_flows.ready
-        yield record_flows
-    finally:
         for server in master.addons.get("proxyserver").servers:
             await server.stop()
         master.shutdown()
+        await task
+
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+
+        # Ignore deprecation warnings raised by mitmproxy at run time.
+        warnings.filterwarnings(
+            "ignore", category=DeprecationWarning, module="mitmproxy"
+        )
+
+        cls.proxy_ready = threading.Event()
+        cls.proxy_thread = threading.Thread(target=asyncio.run, args=(cls.run_proxy(),))
+        cls.proxy_thread.start()
+        cls.proxy_ready.wait()
+
+    def assertNumFlows(self, num_flows):
+        record_flows = self.proxy_master.addons.get("recordflows")
+        self.assertEqual(len(record_flows.get_flows()), num_flows)
 
+    def tearDown(self):
+        record_flows = self.proxy_master.addons.get("recordflows")
+        record_flows.reset_flows()
+        super().tearDown()
 
-@contextlib.contextmanager
-def sync_proxy(mode, **config):
-    loop = None
-    test_done = None
-    proxy_ready = threading.Event()
-    record_flows = None
-
-    async def proxy_coroutine():
-        nonlocal loop, test_done, proxy_ready, record_flows
-        loop = asyncio.get_running_loop()
-        test_done = loop.create_future()
-        async with async_proxy(mode, **config) as record_flows:
-            proxy_ready.set()
-            await test_done
-
-    proxy_thread = threading.Thread(target=asyncio.run, args=(proxy_coroutine(),))
-    proxy_thread.start()
-    try:
-        proxy_ready.wait()
-        yield record_flows
-    finally:
-        loop.call_soon_threadsafe(test_done.set_result, None)
-        proxy_thread.join()
+    @classmethod
+    def tearDownClass(cls):
+        cls.proxy_loop.call_soon_threadsafe(cls.proxy_stop.set_result, None)
+        cls.proxy_thread.join()
+        super().tearDownClass()