Support the Actor class in actor.run()

When actors deriving from `Actor` are passed to `actor.run()`, they will
be started and awaited to finish too.

Signed-off-by: Leandro Lucarella <[email protected]>

Author: llucax

src/frequenz/sdk/actor/_run_utils.py

@@ -8,6 +8,7 @@
 import logging
 from typing import Any
 
+from ._actor import Actor
 from ._decorator import BaseActor
 
 _logger = logging.getLogger(__name__)
@@ -22,14 +23,20 @@ async def run(*actors: Any) -> None:
     Raises:
         AssertionError: if any of the actors is not an instance of BaseActor.
     """
-    # Check that each actor is an instance of BaseActor at runtime,
-    # due to the indirection created by the actor decorator.
-    for actor in actors:
-        assert isinstance(actor, BaseActor), f"{actor} is not an instance of BaseActor"
+    pending_tasks: set[asyncio.Task[Any]] = set()
 
-    pending_tasks = set()
+    # Check that each actor is an instance of BaseActor or Actor at runtime,
+    # due to the indirection created by the actor decorator.
     for actor in actors:
-        pending_tasks.add(asyncio.create_task(actor.join(), name=str(actor)))
+        if isinstance(actor, Actor):
+            await actor.start()
+            awaitable = actor.wait()
+        else:
+            assert isinstance(
+                actor, BaseActor
+            ), f"{actor} is not an instance of BaseActor or Actor"
+            awaitable = actor.join()  # type: ignore
+        pending_tasks.add(asyncio.create_task(awaitable, name=str(actor)))
 
     # Currently the actor decorator manages the life-cycle of the actor tasks
     while pending_tasks:

tests/actor/test_run_utils.py

@@ -11,7 +11,7 @@
 import pytest
 import time_machine
 
-from frequenz.sdk.actor import actor, run
+from frequenz.sdk.actor import Actor, run
 
 
 # Setting 'autouse' has no effect as this method replaces the event loop for all tests in the file.
@@ -30,8 +30,7 @@ def fake_time() -> Iterator[time_machine.Coordinates]:
         yield traveller
 
 
-@actor
-class FaultyActor:
+class FaultyActor(Actor):
     """A test faulty actor."""
 
     def __init__(self, name: str) -> None:
@@ -40,10 +39,10 @@ def __init__(self, name: str) -> None:
         Args:
             name: the name of the faulty actor.
         """
-        self.name = name
+        super().__init__(name=name)
         self.is_cancelled = False
 
-    async def run(self) -> None:
+    async def _run(self) -> None:
         """Run the faulty actor.
 
         Raises:
@@ -53,8 +52,7 @@ async def run(self) -> None:
         raise asyncio.CancelledError(f"Faulty Actor {self.name} failed")
 
 
-@actor
-class SleepyActor:
+class SleepyActor(Actor):
     """A test actor that sleeps a short time."""
 
     def __init__(self, name: str, sleep_duration: float) -> None:
@@ -64,11 +62,11 @@ def __init__(self, name: str, sleep_duration: float) -> None:
             name: the name of the sleepy actor.
             sleep_duration: the virtual duration to sleep while running.
         """
-        self.name = name
+        super().__init__(name=name)
         self.sleep_duration = sleep_duration
         self.is_joined = False
 
-    async def run(self) -> None:
+    async def _run(self) -> None:
         """Run the sleepy actor."""
         while time.time() < self.sleep_duration:
             await asyncio.sleep(0.1)