vucm: explicitly mark tasks and commands

Author: rnovatorov

enapter/vucm/__init__.py

@@ -1,12 +1,14 @@
 from .app import App, run
 from .config import Config
-from .device import Device
+from .device import Device, device_command, device_task
 from .ucm import UCM
 
 __all__ = [
     "App",
     "Config",
     "Device",
+    "device_command",
+    "device_task",
     "UCM",
     "run",
 ]

enapter/vucm/device.py

@@ -2,25 +2,57 @@
 import concurrent
 import functools
 import traceback
-from typing import Optional, Set
+from typing import Any, Callable, Coroutine, Optional, Set
 
 import enapter
 
 from .logger import Logger
 
+DEVICE_TASK_MARK = "_enapter_vucm_device_task"
+DEVICE_COMMAND_MARK = "_enapter_vucm_device_command"
+
+DeviceTaskFunc = Callable[[Any], Coroutine]
+DeviceCommandFunc = Callable[..., Coroutine]
+
+
+def device_task(func: DeviceTaskFunc) -> DeviceTaskFunc:
+    setattr(func, DEVICE_TASK_MARK, True)
+    return func
+
+
+def device_command(func: DeviceCommandFunc) -> DeviceTaskFunc:
+    setattr(func, DEVICE_COMMAND_MARK, True)
+    return func
+
+
+def is_device_task(func: DeviceTaskFunc) -> bool:
+    return getattr(func, DEVICE_TASK_MARK, False) is True
+
+
+def is_device_command(func: DeviceCommandFunc) -> bool:
+    return getattr(func, DEVICE_COMMAND_MARK, False) is True
+
 
 class Device(enapter.async_.Routine):
     def __init__(
         self,
         channel,
         cmd_prefix="cmd_",
-        task_prefix="task_",
         thread_pool_executor=None,
     ) -> None:
         self.__channel = channel
 
-        self.__cmd_prefix = cmd_prefix
-        self.__task_prefix = task_prefix
+        self.__tasks = {}
+        for name in dir(self):
+            obj = getattr(self, name)
+            if is_device_task(obj):
+                self.__tasks[name] = obj
+
+        self.__commands = {}
+        for name in dir(self):
+            obj = getattr(self, name)
+            if is_device_command(obj):
+                self.__commands[name] = obj
 
         if thread_pool_executor is None:
             thread_pool_executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
@@ -62,11 +94,8 @@ async def _run(self):
 
         tasks = set()
 
-        for name in dir(self):
-            if name.startswith(self.__task_prefix):
-                task_func = getattr(self, name)
-                name_without_prefix = name[len(self.__task_prefix) :]
-                tasks.add(asyncio.create_task(task_func(), name=name_without_prefix))
+        for name, func in self.__tasks.items():
+            tasks.add(asyncio.create_task(func(), name=name))
 
         tasks.add(
             asyncio.create_task(
@@ -107,8 +136,8 @@ async def __process_command_requests(self):
 
     async def __execute_command(self, req):
         try:
-            cmd = getattr(self, self.__cmd_prefix + req.name)
-        except AttributeError:
+            cmd = self._commands[req.name]
+        except KeyError:
             return enapter.mqtt.api.CommandState.ERROR, {"reason": "unknown command"}
 
         try:

enapter/vucm/ucm.py

@@ -2,7 +2,7 @@
 
 import enapter
 
-from .device import Device
+from .device import Device, device_command, device_task
 
 
 class UCM(Device):
@@ -13,20 +13,24 @@ def __init__(self, mqtt_client, hardware_id):
             )
         )
 
-    async def cmd_reboot(self):
+    @device_command
+    async def reboot(self):
         await asyncio.sleep(0)
         raise NotImplementedError
 
-    async def cmd_upload_lua_script(self, url, sha1, payload=None):
+    @device_command
+    async def upload_lua_script(self, url, sha1, payload=None):
         await asyncio.sleep(0)
         raise NotImplementedError
 
-    async def task_telemetry_publisher(self):
+    @device_task
+    async def telemetry_publisher(self) -> None:
         while True:
             await self.send_telemetry()
             await asyncio.sleep(1)
 
-    async def task_properties_publisher(self):
+    @device_task
+    async def properties_publisher(self):
         while True:
             await self.send_properties({"virtual": True, "lua_api_ver": 1})
             await asyncio.sleep(10)

tests/unit/test_vucm.py

@@ -11,6 +11,101 @@ async def test_run_in_thread(self, fake):
             assert await device.run_in_thread(lambda: 42) == 42
         assert device._Device__thread_pool_executor._shutdown
 
+    async def test_task_marks(self, fake):
+        class MyDevice(enapter.vucm.Device):
+            async def task_foo(self):
+                pass
+
+            @enapter.vucm.device_task
+            async def task_bar(self):
+                pass
+
+            @enapter.vucm.device_task
+            async def baz(self):
+                pass
+
+            @enapter.vucm.device_task
+            async def goo(self):
+                pass
+
+        async with MyDevice(channel=MockChannel(fake)) as device:
+            tasks = device._Device__tasks
+            assert len(tasks) == 3
+            assert "task_foo" not in tasks
+            assert "task_bar" in tasks
+            assert "baz" in tasks
+            assert "goo" in tasks
+
+    async def test_command_marks(self, fake):
+        class MyDevice(enapter.vucm.Device):
+            async def cmd_foo(self):
+                pass
+
+            async def cmd_foo2(self, a, b, c):
+                pass
+
+            @enapter.vucm.device_command
+            async def cmd_bar(self):
+                pass
+
+            @enapter.vucm.device_command
+            async def cmd_bar2(self, a, b, c):
+                pass
+
+            @enapter.vucm.device_command
+            async def baz(self):
+                pass
+
+            @enapter.vucm.device_command
+            async def baz2(self, a, b, c):
+                pass
+
+            @enapter.vucm.device_command
+            async def goo(self):
+                pass
+
+        async with MyDevice(channel=MockChannel(fake)) as device:
+            commands = device._Device__commands
+            assert len(commands) == 5
+            assert "cmd_foo" not in commands
+            assert "cmd_foo2" not in commands
+            assert "cmd_bar" in commands
+            assert "cmd_bar2" in commands
+            assert "baz" in commands
+            assert "baz2" in commands
+            assert "goo" in commands
+
+    async def test_task_and_commands_marks(self, fake):
+        class MyDevice(enapter.vucm.Device):
+            @enapter.vucm.device_task
+            async def foo_task(self):
+                pass
+
+            @enapter.vucm.device_task
+            async def bar_task(self):
+                pass
+
+            @enapter.vucm.device_command
+            async def foo_command(self):
+                pass
+
+            @enapter.vucm.device_command
+            async def bar_command(self):
+                pass
+
+        async with MyDevice(channel=MockChannel(fake)) as device:
+            tasks = device._Device__tasks
+            assert len(tasks) == 2
+            assert "foo_task" in tasks
+            assert "bar_task" in tasks
+            assert "foo_command" not in tasks
+            assert "bar_command" not in tasks
+            commands = device._Device__commands
+            assert "foo_task" not in commands
+            assert "bar_task" not in commands
+            assert "foo_command" in commands
+            assert "bar_command" in commands
+
 
 class MockChannel:
     def __init__(self, fake):