additional unit tests

src/asyncio_taskpool/exceptions.py

@@ -55,5 +55,13 @@ class ServerException(Exception):
     pass
 
 
+class UnknownTaskPoolClass(ServerException):
+    pass
+
+
+class NotATaskPool(ServerException):
+    pass
+
+
 class HelpRequested(ServerException):
     pass

src/asyncio_taskpool/helpers.py

@@ -19,7 +19,6 @@ Miscellaneous helper functions.
 """
 
 
-import re
 from asyncio.coroutines import iscoroutinefunction
 from asyncio.queues import Queue
 from inspect import getdoc
@@ -57,7 +56,7 @@ def tasks_str(num: int) -> str:
 
 
 def get_first_doc_line(obj: object) -> str:
-    return getdoc(obj).strip().split("\n", 1)[0]
+    return getdoc(obj).strip().split("\n", 1)[0].strip()
 
 
 async def return_or_exception(_function_to_execute: AnyCallableT, *args, **kwargs) -> Union[T, Exception]:

src/asyncio_taskpool/session.py

@@ -5,9 +5,9 @@ from asyncio.streams import StreamReader, StreamWriter
 from typing import Callable, Optional, Union, TYPE_CHECKING
 
 from .constants import CMD, SESSION_PARSER_WRITER, SESSION_MSG_BYTES, CLIENT_INFO
-from .exceptions import HelpRequested
+from .exceptions import HelpRequested, NotATaskPool, UnknownTaskPoolClass
 from .helpers import get_first_doc_line, return_or_exception, tasks_str
-from .pool import TaskPool, SimpleTaskPool
+from .pool import BaseTaskPool, TaskPool, SimpleTaskPool
 from .session_parser import CommandParser, NUM
 
 if TYPE_CHECKING:
@@ -67,6 +67,9 @@ class ControlSession:
             CMD.FUNC_NAME, short_help=get_first_doc_line(self._pool.__class__.func_name.fget)
         )
 
+    def _add_advanced_commands(self) -> None:
+        raise NotImplementedError
+
     def _init_parser(self, client_terminal_width: int) -> None:
         parser_kwargs = {
             'prog': '',
@@ -76,9 +79,13 @@ class ControlSession:
         self._parser = CommandParser(**parser_kwargs)
         self._add_base_commands()
         if isinstance(self._pool, TaskPool):
-            pass  # TODO
+            self._add_advanced_commands()
         elif isinstance(self._pool, SimpleTaskPool):
             self._add_simple_commands()
+        elif isinstance(self._pool, BaseTaskPool):
+            raise UnknownTaskPoolClass(f"No interface defined for {self._pool.__class__.__name__}")
+        else:
+            raise NotATaskPool(f"Not a task pool instance: {self._pool}")
 
     async def client_handshake(self) -> None:
         client_info = json.loads((await self._reader.read(SESSION_MSG_BYTES)).decode().strip())

tests/test_helpers.py

@@ -94,3 +94,35 @@ class HelpersTestCase(IsolatedAsyncioTestCase):
         self.assertEqual("tasks", helpers.tasks_str(2))
         self.assertEqual("tasks", helpers.tasks_str(-10))
         self.assertEqual("tasks", helpers.tasks_str(42))
+
+    def test_get_first_doc_line(self):
+        expected_output = 'foo bar baz'
+        mock_obj = MagicMock(__doc__=f"""{expected_output} 
+        something else
+        
+        even more
+        """)
+        output = helpers.get_first_doc_line(mock_obj)
+        self.assertEqual(expected_output, output)
+
+    async def test_return_or_exception(self):
+        expected_output = '420'
+        mock_func = AsyncMock(return_value=expected_output)
+        args = (1, 3, 5)
+        kwargs = {'a': 1, 'b': 2, 'c': 'foo'}
+        output = await helpers.return_or_exception(mock_func, *args, **kwargs)
+        self.assertEqual(expected_output, output)
+        mock_func.assert_awaited_once_with(*args, **kwargs)
+
+        mock_func = MagicMock(return_value=expected_output)
+        output = await helpers.return_or_exception(mock_func, *args, **kwargs)
+        self.assertEqual(expected_output, output)
+        mock_func.assert_called_once_with(*args, **kwargs)
+
+        class TestException(Exception):
+            pass
+        test_exception = TestException()
+        mock_func = MagicMock(side_effect=test_exception)
+        output = await helpers.return_or_exception(mock_func, *args, **kwargs)
+        self.assertEqual(test_exception, output)
+        mock_func.assert_called_once_with(*args, **kwargs)