[Coverage] For newly added components (#1014)

* Add newly added code cov

* Fix spelling

Author: abhinavsingh

examples/ssl_echo_server.py

@@ -21,7 +21,7 @@ class EchoSSLServerHandler(BaseTcpServerHandler[TcpClientConnection]):
     """Wraps client socket during initialization."""
 
     @staticmethod
-    def create(**kwargs: Any) -> TcpClientConnection:
+    def create(**kwargs: Any) -> TcpClientConnection:   # pragma: no cover
         return TcpClientConnection(**kwargs)
 
     def initialize(self) -> None:

examples/tcp_echo_server.py

@@ -20,7 +20,7 @@ class EchoServerHandler(BaseTcpServerHandler[TcpClientConnection]):
     """Sets client socket to non-blocking during initialization."""
 
     @staticmethod
-    def create(**kwargs: Any) -> TcpClientConnection:
+    def create(**kwargs: Any) -> TcpClientConnection:   # pragma: no cover
         return TcpClientConnection(**kwargs)
 
     def initialize(self) -> None:

proxy/core/base/tcp_tunnel.py

@@ -48,7 +48,7 @@ def handle_data(self, data: memoryview) -> Optional[bool]:
         pass    # pragma: no cover
 
     @staticmethod
-    def create(**kwargs: Any) -> TcpClientConnection:
+    def create(**kwargs: Any) -> TcpClientConnection:   # pragma: no cover
         return TcpClientConnection(**kwargs)
 
     def initialize(self) -> None:

proxy/core/connection/pool.py

@@ -78,7 +78,7 @@ def __init__(self) -> None:
         self.pools: Dict[Tuple[str, int], Set[TcpServerConnection]] = {}
 
     @staticmethod
-    def create(**kwargs: Any) -> TcpServerConnection:
+    def create(**kwargs: Any) -> TcpServerConnection:   # pragma: no cover
         return TcpServerConnection(**kwargs)
 
     def acquire(self, addr: Tuple[str, int]) -> Tuple[bool, TcpServerConnection]:

proxy/core/ssh/listener.py

@@ -18,7 +18,7 @@
     from paramiko.transport import Transport
     if TYPE_CHECKING:   # pragma: no cover
         from paramiko.channel import Channel
-except ImportError:
+except ImportError:     # pragma: no cover
     pass
 
 from ...common.flag import flags

proxy/http/handler.py

@@ -58,7 +58,7 @@ def __init__(self, *args: Any, **kwargs: Any):
     ##
 
     @staticmethod
-    def create(**kwargs: Any) -> HttpClientConnection:
+    def create(**kwargs: Any) -> HttpClientConnection:  # pragma: no cover
         return HttpClientConnection(**kwargs)
 
     def initialize(self) -> None:
@@ -205,12 +205,14 @@ async def handle_writables(self, writables: Writables) -> bool:
                 if teardown:
                     return True
             except BrokenPipeError:
-                logger.warning(
+                logger.warning(     # pragma: no cover
                     'BrokenPipeError when flushing buffer for client',
                 )
                 return True
             except OSError:
-                logger.warning('OSError when flushing buffer to client')
+                logger.warning(     # pragma: no cover
+                    'OSError when flushing buffer to client',
+                )
                 return True
         return False
 
@@ -261,8 +263,6 @@ def _initialize_plugin(
 
     def _discover_plugin_klass(self, protocol: int) -> Optional[Type['HttpProtocolHandlerPlugin']]:
         """Discovers and return matching HTTP handler plugin matching protocol."""
-        if protocol == httpProtocols.UNKNOWN:
-            return None
         if b'HttpProtocolHandlerPlugin' in self.flags.plugins:
             for klass in self.flags.plugins[b'HttpProtocolHandlerPlugin']:
                 k: Type['HttpProtocolHandlerPlugin'] = klass
@@ -287,6 +287,10 @@ def _parse_first_request(self, data: memoryview) -> bool:
             ) from e
         if not self.request.is_complete:
             return False
+        # Bail out if http protocol is unknown
+        if self.request.http_handler_protocol == httpProtocols.UNKNOWN:
+            self.work.queue(BAD_REQUEST_RESPONSE_PKT)
+            return True
         # Discover which HTTP handler plugin is capable of
         # handling the current incoming request
         klass = self._discover_plugin_klass(

tests/http/parser/test_http_parser.py

@@ -17,6 +17,7 @@
 )
 from proxy.http.exception import HttpProtocolException
 from proxy.http.responses import okResponse
+from proxy.http.protocols import httpProtocols
 from proxy.common.constants import CRLF, HTTP_1_0
 
 
@@ -832,6 +833,10 @@ def test_parses_icap_protocol(self) -> None:
         self.assertEqual(self.parser.method, b'REQMOD')
         assert self.parser._url is not None
         self.assertEqual(self.parser._url.scheme, b'icap')
+        self.assertEqual(
+            self.parser.http_handler_protocol,
+            httpProtocols.UNKNOWN,
+        )
 
     def test_cannot_parse_sip_protocol(self) -> None:
         # Will fail to parse because of invalid host and port in the request line

tests/test_main.py

@@ -308,11 +308,51 @@ def test_enable_devtools(
         mock_acceptor_pool.return_value.setup.assert_called_once()
         mock_listener.return_value.setup.assert_called_once()
 
-    # def test_pac_file(self) -> None:
-    #     pass
-
-    # def test_imports_plugin(self) -> None:
-    #     pass
-
-    # def test_cannot_enable_https_proxy_and_tls_interception_mutually(self) -> None:
-    #     pass
+    @mock.patch('time.sleep')
+    @mock.patch('proxy.common.plugins.Plugins.load')
+    @mock.patch('proxy.common.flag.FlagParser.parse_args')
+    @mock.patch('proxy.proxy.EventManager')
+    @mock.patch('proxy.proxy.AcceptorPool')
+    @mock.patch('proxy.proxy.ThreadlessPool')
+    @mock.patch('proxy.proxy.Listener')
+    @mock.patch('proxy.proxy.SshHttpProtocolHandler')
+    @mock.patch('proxy.proxy.SshTunnelListener')
+    def test_enable_ssh_tunnel(
+            self,
+            mock_ssh_tunnel_listener: mock.Mock,
+            mock_ssh_http_proto_handler: mock.Mock,
+            mock_listener: mock.Mock,
+            mock_executor_pool: mock.Mock,
+            mock_acceptor_pool: mock.Mock,
+            mock_event_manager: mock.Mock,
+            mock_parse_args: mock.Mock,
+            mock_load_plugins: mock.Mock,
+            mock_sleep: mock.Mock,
+    ) -> None:
+        mock_sleep.side_effect = KeyboardInterrupt()
+        mock_args = mock_parse_args.return_value
+        self.mock_default_args(mock_args)
+        mock_args.enable_ssh_tunnel = True
+        mock_args.local_executor = 0
+        main(enable_ssh_tunnel=True, local_executor=0)
+        mock_load_plugins.assert_called()
+        self.assertEqual(
+            mock_load_plugins.call_args_list[0][0][0], [
+                bytes_(PLUGIN_HTTP_PROXY),
+            ],
+        )
+        mock_parse_args.assert_called_once()
+        mock_event_manager.assert_not_called()
+        if _env_threadless_compliant():
+            mock_executor_pool.assert_called_once()
+            mock_executor_pool.return_value.setup.assert_called_once()
+        mock_acceptor_pool.assert_called_once()
+        mock_acceptor_pool.return_value.setup.assert_called_once()
+        mock_listener.return_value.setup.assert_called_once()
+        mock_ssh_http_proto_handler.assert_called_once()
+        mock_ssh_tunnel_listener.assert_called_once()
+        mock_ssh_tunnel_listener.return_value.setup.assert_called_once()
+        mock_ssh_tunnel_listener.return_value.start_port_forward.assert_called_once()
+        mock_ssh_tunnel_listener.return_value.shutdown.assert_called_once()
+        # shutdown will internally call stop port forward
+        mock_ssh_tunnel_listener.return_value.stop_port_forward.assert_not_called()