Fix tests

proxy/common/flag.py

@@ -84,8 +84,8 @@ def parse_args(
 
     @staticmethod
     def initialize(
-        input_args: Optional[List[str]]
-        = None, **opts: Any,
+        input_args: Optional[List[str]] = None,
+        **opts: Any,
     ) -> argparse.Namespace:
         if input_args is None:
             input_args = []
@@ -126,9 +126,9 @@ def initialize(
 
         # Load work_klass
         work_klass = opts.get('work_klass', args.work_klass)
-        work_klass = work_klass \
-            if isinstance(work_klass, type) \
-            else Plugins.importer(work_klass)[0]
+        work_klass = Plugins.importer(work_klass)[0] \
+            if isinstance(work_klass, str) \
+            else work_klass
 
         # Generate auth_code required for basic authentication if enabled
         auth_code = None

proxy/core/acceptor/acceptor.py

@@ -18,12 +18,10 @@
 from multiprocessing import connection
 from multiprocessing.reduction import recv_handle
 
-from typing import List, Optional, Tuple, Type
+from typing import List, Optional, Tuple
 
 from proxy.core.acceptor.executors import ThreadlessPool
 
-from .work import Work
-
 from ..event import EventQueue
 
 from ...common.utils import is_threadless

proxy/core/acceptor/executors.py

@@ -17,7 +17,7 @@
 from multiprocessing import connection
 from multiprocessing.reduction import send_handle
 
-from typing import Optional, List, Tuple, Type
+from typing import Optional, List, Tuple
 
 from .work import Work
 from .threadless import Threadless

proxy/proxy.py

@@ -12,11 +12,10 @@
 import time
 import logging
 
-from typing import List, Optional, Any, Type
+from typing import List, Optional, Any
 
-from .core.acceptor import AcceptorPool, ThreadlessPool, Work
+from .core.acceptor import AcceptorPool, ThreadlessPool
 from .core.event import EventManager
-from .http import HttpProtocolHandler
 from .common.flag import FlagParser, flags
 from .common.constants import DEFAULT_LOG_FILE, DEFAULT_LOG_FORMAT, DEFAULT_LOG_LEVEL
 from .common.constants import DEFAULT_OPEN_FILE_LIMIT, DEFAULT_PLUGINS, DEFAULT_VERSION
@@ -107,12 +106,8 @@ class Proxy(SetupShutdownContextManager):
     for message sharing and/or signaling.
     """
 
-    def __init__(self, **opts: Any) -> None:
-        input_args = sys.argv[1:]
-        print(input_args)
-        print('*'*20)
+    def __init__(self, input_args: Optional[List[str]] = None, **opts: Any) -> None:
         self.flags = FlagParser.initialize(input_args, **opts)
-        print(self.flags)
         self.acceptors: Optional[AcceptorPool] = None
         self.executors: Optional[ThreadlessPool] = None
         self.event_manager: Optional[EventManager] = None
@@ -174,7 +169,7 @@ def shutdown(self) -> None:
 
 def main(**opts: Any) -> None:
     try:
-        with Proxy(**opts):
+        with Proxy(sys.argv[1:], **opts):
             while True:
                 time.sleep(1)
     except KeyboardInterrupt:

proxy/testing/test_case.py

@@ -41,7 +41,7 @@ def setUpClass(cls) -> None:
         cls.INPUT_ARGS.append('--port')
         cls.INPUT_ARGS.append('0')
 
-        cls.PROXY = Proxy(input_args=cls.INPUT_ARGS)
+        cls.PROXY = Proxy(cls.INPUT_ARGS)
         cls.PROXY.flags.plugins[b'HttpProxyBasePlugin'].append(
             CacheResponsesPlugin,
         )

tests/common/test_flags.py

@@ -8,16 +8,18 @@
     :copyright: (c) 2013-present by Abhinav Singh and contributors.
     :license: BSD, see LICENSE for more details.
 """
-from proxy.common.utils import bytes_
-from proxy.common.constants import PLUGIN_HTTP_PROXY
 import unittest
 
+from unittest import mock
 from typing import List, Dict
 
-from proxy.common.flag import FlagParser
 from proxy.http.proxy import HttpProxyPlugin
 from proxy.plugin import CacheResponsesPlugin
 from proxy.plugin import FilterByUpstreamHostPlugin
+from proxy.common.utils import bytes_
+from proxy.common.flag import FlagParser
+from proxy.common.version import __version__
+from proxy.common.constants import PLUGIN_HTTP_PROXY, PY2_DEPRECATION_MESSAGE
 
 
 class TestFlags(unittest.TestCase):
@@ -139,6 +141,43 @@ def test_unique_plugin_from_class(self) -> None:
             ],
         })
 
+    def test_basic_auth_flag_is_base64_encoded(self) -> None:
+        flags = FlagParser.initialize(['--basic-auth', 'user:pass'])
+        self.assertEqual(flags.auth_code, b'dXNlcjpwYXNz')
+
+    @mock.patch('builtins.print')
+    def test_main_version(self, mock_print: mock.Mock) -> None:
+        with self.assertRaises(SystemExit) as e:
+            FlagParser.initialize(['--version'])
+            mock_print.assert_called_with(__version__)
+        self.assertEqual(e.exception.code, 0)
+
+    @mock.patch('builtins.print')
+    @mock.patch('proxy.common.flag.is_py2')
+    def test_main_py2_exit(
+            self,
+            mock_is_py2: mock.Mock,
+            mock_print: mock.Mock,
+    ) -> None:
+        mock_is_py2.return_value = True
+        with self.assertRaises(SystemExit) as e:
+            FlagParser.initialize()
+        mock_print.assert_called_with(PY2_DEPRECATION_MESSAGE)
+        self.assertEqual(e.exception.code, 1)
+        mock_is_py2.assert_called()
+
+    @mock.patch('builtins.print')
+    @mock.patch('proxy.common.flag.is_py2')
+    def test_main_py3_runs(
+            self,
+            mock_is_py2: mock.Mock,
+            mock_print: mock.Mock,
+    ) -> None:
+        mock_is_py2.return_value = False
+        FlagParser.initialize()
+        mock_is_py2.assert_called()
+        mock_print.assert_not_called()
+
 
 if __name__ == '__main__':
     unittest.main()

tests/core/test_acceptor.py

@@ -16,21 +16,23 @@
 
 from proxy.core.acceptor import Acceptor
 from proxy.common.flag import FlagParser
+from proxy.http import HttpProtocolHandler
 
 
 class TestAcceptor(unittest.TestCase):
 
     def setUp(self) -> None:
         self.acceptor_id = 1
-        self.mock_protocol_handler = mock.MagicMock()
         self.pipe = multiprocessing.Pipe()
-        self.flags = FlagParser.initialize(threaded=True)
+        self.flags = FlagParser.initialize(
+            threaded=True, work_klass=mock.MagicMock())
         self.acceptor = Acceptor(
             idd=self.acceptor_id,
             work_queue=self.pipe[1],
             flags=self.flags,
             lock=multiprocessing.Lock(),
-            work_klass=self.mock_protocol_handler,
+            executor_queues=[],
+            executor_pids=[],
         )
 
     @mock.patch('selectors.DefaultSelector')
@@ -55,9 +57,9 @@ def test_continues_when_no_events(
         self.acceptor.run()
 
         sock.accept.assert_not_called()
-        self.mock_protocol_handler.assert_not_called()
+        self.flags.work_klass.assert_not_called()
 
-    @mock.patch('proxy.core.acceptor.acceptor.TcpClientConnection')
+    @mock.patch('proxy.core.acceptor.executors.TcpClientConnection')
     @mock.patch('threading.Thread')
     @mock.patch('selectors.DefaultSelector')
     @mock.patch('socket.fromfd')
@@ -92,13 +94,13 @@ def test_accepts_client_from_server_socket(
             family=socket.AF_INET6,
             type=socket.SOCK_STREAM,
         )
-        self.mock_protocol_handler.assert_called_with(
+        self.flags.work_klass.assert_called_with(
             mock_client.return_value,
             flags=self.flags,
             event_queue=None,
         )
         mock_thread.assert_called_with(
-            target=self.mock_protocol_handler.return_value.run,
+            target=self.flags.work_klass.return_value.run,
         )
         mock_thread.return_value.start.assert_called()
         sock.close.assert_called()

tests/core/test_acceptor_pool.py

@@ -46,16 +46,14 @@ def test_setup_and_shutdown(
         num_workers = 2
         pid_file = os.path.join(tempfile.gettempdir(), 'pid')
         sock = mock_socket.return_value
-        work_klass = mock.MagicMock()
         flags = FlagParser.initialize(
             num_workers=2, pid_file=pid_file, threaded=True,
         )
 
-        pool = AcceptorPool(flags=flags, work_klass=work_klass)
+        pool = AcceptorPool(flags=flags, executor_queues=[], executor_pids=[])
         pool.setup()
         mock_send_handle.assert_called()
 
-        work_klass.assert_not_called()
         mock_socket.assert_called_with(
             socket.AF_INET6 if pool.flags.hostname.version == 6 else socket.AF_INET,
             socket.SOCK_STREAM,

tests/test_main.py

@@ -14,8 +14,6 @@
 
 from proxy.proxy import main, entry_point
 from proxy.common.utils import bytes_
-from proxy.http import HttpProtocolHandler
-from proxy.common.flag import FlagParser
 
 from proxy.common.constants import DEFAULT_ENABLE_DASHBOARD, DEFAULT_LOG_LEVEL, DEFAULT_LOG_FILE, DEFAULT_LOG_FORMAT, DEFAULT_NUM_ACCEPTORS, DEFAULT_WORK_KLASS
 from proxy.common.constants import DEFAULT_TIMEOUT, DEFAULT_DEVTOOLS_WS_PATH, DEFAULT_DISABLE_HTTP_PROXY
@@ -240,82 +238,6 @@ def test_enable_devtools(
         mock_acceptor_pool.assert_called_once()
         mock_acceptor_pool.return_value.setup.assert_called_once()
 
-    @mock.patch('time.sleep')
-    @mock.patch('proxy.proxy.EventManager')
-    @mock.patch('proxy.proxy.AcceptorPool')
-    def test_basic_auth_flag_is_base64_encoded(
-            self,
-            mock_acceptor_pool: mock.Mock,
-            mock_event_manager: mock.Mock,
-            mock_sleep: mock.Mock,
-    ) -> None:
-        mock_sleep.side_effect = KeyboardInterrupt()
-
-        input_args = ['--basic-auth', 'user:pass']
-        flgs = FlagParser.initialize(input_args)
-
-        main(input_args=input_args)
-        mock_event_manager.assert_not_called()
-        mock_acceptor_pool.assert_called_once()
-        self.assertEqual(
-            flgs.auth_code,
-            b'dXNlcjpwYXNz',
-        )
-
-    @mock.patch('time.sleep')
-    @mock.patch('builtins.print')
-    @mock.patch('proxy.proxy.EventManager')
-    @mock.patch('proxy.proxy.AcceptorPool')
-    @mock.patch('proxy.common.flag.is_py2')
-    def test_main_py3_runs(
-            self,
-            mock_is_py2: mock.Mock,
-            mock_acceptor_pool: mock.Mock,
-            mock_event_manager: mock.Mock,
-            mock_print: mock.Mock,
-            mock_sleep: mock.Mock,
-    ) -> None:
-        mock_sleep.side_effect = KeyboardInterrupt()
-
-        input_args = ['--basic-auth', 'user:pass']
-        mock_is_py2.return_value = False
-
-        main(input_args, num_workers=1)
-
-        mock_is_py2.assert_called()
-        mock_print.assert_not_called()
-
-        mock_event_manager.assert_not_called()
-        mock_acceptor_pool.assert_called_once()
-        mock_acceptor_pool.return_value.setup.assert_called()
-
-    @mock.patch('builtins.print')
-    @mock.patch('proxy.common.flag.is_py2')
-    def test_main_py2_exit(
-            self,
-            mock_is_py2: mock.Mock,
-            mock_print: mock.Mock,
-    ) -> None:
-        mock_is_py2.return_value = True
-        with self.assertRaises(SystemExit) as e:
-            main(num_workers=1)
-        mock_print.assert_called_with(PY2_DEPRECATION_MESSAGE)
-        self.assertEqual(e.exception.code, 1)
-        mock_is_py2.assert_called()
-
-    @mock.patch('builtins.print')
-    @mock.patch('sys.argv')
-    def test_main_version(
-            self,
-            mock_sys_argv: mock.Mock,
-            mock_print: mock.Mock,
-    ) -> None:
-        mock_sys_argv.return_value = ['proxy', '--version']
-        with self.assertRaises(SystemExit) as e:
-            main()
-            mock_print.assert_called_with(__version__)
-        self.assertEqual(e.exception.code, 0)
-
     # def test_pac_file(self) -> None:
     #     pass