finally reached 100% unittest coverage overall

src/asyncio_taskpool/control/__main__.py

@@ -19,11 +19,10 @@ CLI client entry point.
 """
 
 
-import sys
 from argparse import ArgumentParser
 from asyncio import run
 from pathlib import Path
-from typing import Dict, Any
+from typing import Any, Dict, Sequence
 
 from ..constants import PACKAGE_NAME
 from ..pool import TaskPool
@@ -31,25 +30,19 @@ from .client import ControlClient, TCPControlClient, UnixControlClient
 from .server import TCPControlServer, UnixControlServer
 
 
-CONN_TYPE = 'conn_type'
+CLIENT_CLASS = 'client_class'
 UNIX, TCP = 'unix', 'tcp'
 SOCKET_PATH = 'path'
 HOST, PORT = 'host', 'port'
 
 
-def parse_cli() -> Dict[str, Any]:
+def parse_cli(args: Sequence[str] = None) -> Dict[str, Any]:
     parser = ArgumentParser(
-        prog=PACKAGE_NAME,
-        description=f"CLI based {ControlClient.__name__} for {PACKAGE_NAME}"
-    )
-    subparsers = parser.add_subparsers(title="Connection types", dest=CONN_TYPE)
-    unix_parser = subparsers.add_parser(UNIX, help="Connect via unix socket")
-    unix_parser.add_argument(
-        SOCKET_PATH,
-        type=Path,
-        help=f"Path to the unix socket on which the {UnixControlServer.__name__} for the {TaskPool.__name__} is "
-             f"listening."
+        prog=f'{PACKAGE_NAME}.control',
+        description=f"Simple CLI based {ControlClient.__name__} for {PACKAGE_NAME}"
     )
+    subparsers = parser.add_subparsers(title="Connection types")
+
     tcp_parser = subparsers.add_parser(TCP, help="Connect via TCP socket")
     tcp_parser.add_argument(
         HOST,
@@ -60,19 +53,25 @@ def parse_cli() -> Dict[str, Any]:
         type=int,
         help=f"Port that the {TCPControlServer.__name__} for the {TaskPool.__name__} is listening on."
     )
-    return vars(parser.parse_args())
+    tcp_parser.set_defaults(**{CLIENT_CLASS: TCPControlClient})
+
+    unix_parser = subparsers.add_parser(UNIX, help="Connect via unix socket")
+    unix_parser.add_argument(
+        SOCKET_PATH,
+        type=Path,
+        help=f"Path to the unix socket on which the {UnixControlServer.__name__} for the {TaskPool.__name__} is "
+             f"listening."
+    )
+    unix_parser.set_defaults(**{CLIENT_CLASS: UnixControlClient})
+
+    return vars(parser.parse_args(args))
 
 
 async def main():
     kwargs = parse_cli()
-    if kwargs[CONN_TYPE] == UNIX:
-        client = UnixControlClient(socket_path=kwargs[SOCKET_PATH])
-    elif kwargs[CONN_TYPE] == TCP:
-        client = TCPControlClient(host=kwargs[HOST], port=kwargs[PORT])
-    else:
-        print("Invalid connection type", file=sys.stderr)
-        sys.exit(2)
-    await client.start()
+    client_cls = kwargs.pop(CLIENT_CLASS)
+    await client_cls(**kwargs).start()
+
 
 if __name__ == '__main__':
     run(main())