Define a way for constructing objects with async constructor

Add parent class that raises error if __init__ was called.
Use this for implementing BatteryPoolStatus object.

Signed-off-by: ela-kotulska-frequenz <[email protected]>

Author: ela-kotulska-frequenz

src/frequenz/sdk/_internal/asyncio.py

@@ -4,6 +4,7 @@
 """General purpose async tools."""
 
 import asyncio
+from abc import ABC
 
 
 async def cancel_and_await(task: asyncio.Task) -> None:
@@ -19,3 +20,22 @@ async def cancel_and_await(task: asyncio.Task) -> None:
         await task
     except asyncio.CancelledError:
         pass
+
+
+class NotSyncConstructible(AssertionError):
+    """Raised when object with async constructor is created in sync way."""
+
+
+class AsyncConstructible(ABC):
+    """Parent class for classes where part of the constructor is async."""
+
+    def __init__(self) -> None:
+        """Raise error when object is created in sync way.
+
+        Raises:
+            NotSyncConstructible: If this method is called.
+        """
+        raise NotSyncConstructible(
+            "This object shouldn't be created with default constructor. ",
+            "Check class documentation for more information.",
+        )

src/frequenz/sdk/actor/power_distributing/_battery_pool_status.py

@@ -2,30 +2,39 @@
 # Copyright © 2022 Frequenz Energy-as-a-Service GmbH
 """Class that stores pool of batteries and manage them."""
 
+from __future__ import annotations
+
 import asyncio
 import logging
-from typing import Set
+from typing import Dict, Set
 
+from ..._internal.asyncio import AsyncConstructible
 from ...microgrid._battery import BatteryStatus, StatusTracker
 from .result import PartialFailure, Result, Success
 
 _logger = logging.getLogger(__name__)
 
 
-class BatteryPoolStatus:
-    """Holds pool of batteries and returns data from them."""
+class BatteryPoolStatus(AsyncConstructible):
+    """Return status of batteries in the pool.
+
+    To create an instance of this class you should use `async_new` class method.
+    Standard constructor (__init__) is not supported and using it will raise
+    `NotSyncConstructible` error.
+    """
 
-    def __init__(
-        self,
+    # This is instance attribute.
+    # Don't assign default value, because then it becomes class attribute.
+    _batteries: Dict[int, StatusTracker]
+
+    @classmethod
+    async def async_new(
+        cls,
         battery_ids: Set[int],
         max_data_age_sec: float,
         max_blocking_duration_sec: float,
-    ) -> None:
-        """Create partially initialized object instance.
-
-        Note:
-            Please call `async_init` method to fully initialize BatteryPoolStatus. Otherwise
-            it is not possible to use BatteryPoolStatus.
+    ) -> BatteryPoolStatus:
+        """Create BatteryPoolStatus instance.
 
         Args:
             battery_ids: set of batteries ids that should be stored in pool.
@@ -35,20 +44,21 @@ def __init__(
                 data.
             max_blocking_duration_sec: This value tell what should be the maximum
                 timeout used for blocking failing component.
+
+        Returns:
+            New instance of this class.
         """
+        self: BatteryPoolStatus = BatteryPoolStatus.__new__(cls)
         self._batteries = {
             id: StatusTracker(id, max_data_age_sec, max_blocking_duration_sec)
             for id in battery_ids
         }
-        self._init_method_called: bool = False
 
-    async def async_init(self) -> None:
-        """Init battery pool."""
         await asyncio.gather(
             *[bat.async_init() for bat in self._batteries.values()],
             return_exceptions=True,
         )
-        self._init_method_called = True
+        return self
 
     def get_working_batteries(self, battery_ids: Set[int]) -> Set[int]:
         """Get subset of battery_ids with working batteries.
@@ -64,11 +74,6 @@ def get_working_batteries(self, battery_ids: Set[int]) -> Set[int]:
         Returns:
             Subset of given batteries with working batteries.
         """
-        if not self._init_method_called:
-            raise RuntimeError(
-                "`async_init` method not called or not awaited. Run it before first use"
-            )
-
         working: Set[int] = set()
         uncertain: Set[int] = set()
         for bat_id in battery_ids:
@@ -101,11 +106,6 @@ def update_last_request_status(self, result: Result):
             RuntimeError: If `async_init` method was not called at the beginning to
                 initialize object.
         """
-        if not self._init_method_called:
-            raise RuntimeError(
-                "`async_init` method not called or not awaited. Run it before first use"
-            )
-
         if isinstance(result, Success):
             for bat_id in result.used_batteries:
                 self._batteries[bat_id].unblock()

src/frequenz/sdk/actor/power_distributing/power_distributing.py

@@ -161,16 +161,11 @@ def __init__(
             self.power_distributor_exponent
         )
 
-        graph = microgrid.get().component_graph
-        batteries = graph.components(component_category={ComponentCategory.BATTERY})
+        self._battery_pool: BatteryPoolStatus
 
-        self._battery_pool = BatteryPoolStatus(
-            battery_ids={battery.component_id for battery in batteries},
-            max_blocking_duration_sec=30.0,
-            max_data_age_sec=10.0,
+        self._bat_inv_map, self._inv_bat_map = self._get_components_pairs(
+            microgrid.get().component_graph
         )
-
-        self._bat_inv_map, self._inv_bat_map = self._get_components_pairs(graph)
         self._battery_receivers: Dict[int, Peekable[BatteryData]] = {}
         self._inverter_receivers: Dict[int, Peekable[InverterData]] = {}
 
@@ -251,11 +246,17 @@ async def run(self) -> None:
         as broken for some time.
         """
         await self._create_channels()
-        await self._battery_pool.async_init()
+
         api = microgrid.get().api_client
+        self._battery_pool = await BatteryPoolStatus.async_new(
+            battery_ids=set(self._bat_inv_map.keys()),
+            max_blocking_duration_sec=30.0,
+            max_data_age_sec=10.0,
+        )
 
         # Wait few seconds to get data from the channels created above.
         await asyncio.sleep(self._wait_for_data_sec)
+
         self._started.set()
         while True:
             request, user = await self._request_queue.get()

tests/actor/test_battery_pool.py

@@ -7,6 +7,7 @@
 from typing import Callable, Iterable, List, Optional, Set
 
 import frequenz.api.microgrid.microgrid_pb2 as microgrid_pb
+import pytest
 import pytz
 import time_machine
 from frequenz.api.microgrid.battery_pb2 import Battery as PbBattery
@@ -25,6 +26,7 @@
 from pytest_mock import MockerFixture
 
 from frequenz.sdk import microgrid
+from frequenz.sdk._internal.asyncio import NotSyncConstructible
 from frequenz.sdk.actor.power_distributing import PartialFailure, Request, Success
 from frequenz.sdk.actor.power_distributing._battery_pool_status import BatteryPoolStatus
 from frequenz.sdk.microgrid.component import BatteryData, InverterData
@@ -453,6 +455,11 @@ def set_battery_error(
             ),
         )
 
+    def test_create_sync_pool(self) -> None:
+        """Test if error is raised after calling default constructor."""
+        with pytest.raises(NotSyncConstructible):
+            BatteryPoolStatus()
+
     async def test_scenario_one_battery(self, mocker: MockerFixture) -> None:
         """Test scenario with one battery.
 
@@ -462,10 +469,9 @@ async def test_scenario_one_battery(self, mocker: MockerFixture) -> None:
         pairs = await self.my_setup(mocker=mocker, batteries_num=1)
 
         batteries = {pair.bat_id for pair in pairs}
-        battery_pool = BatteryPoolStatus(
+        battery_pool = await BatteryPoolStatus.async_new(
             batteries, max_data_age_sec=30, max_blocking_duration_sec=30
         )
-        await battery_pool.async_init()
 
         self.scenario_all_msg_correct(battery_pool, pairs)
         self.scenario_no_message(battery_pool, pairs)
@@ -484,10 +490,9 @@ async def test_scenario_five_batteries(self, mocker: MockerFixture) -> None:
         """
         pairs = await self.my_setup(mocker=mocker, batteries_num=5)
         batteries = {pair.bat_id for pair in pairs}
-        battery_pool = BatteryPoolStatus(
+        battery_pool = await BatteryPoolStatus.async_new(
             batteries, max_data_age_sec=30, max_blocking_duration_sec=30
         )
-        await battery_pool.async_init()
 
         # Starting server takes long time, so it is better to start it once and
         # then run tests.

tests/actor/test_power_distributing.py

@@ -138,13 +138,15 @@ async def test_power_distributor_one_user(self, mocker: MockerFixture) -> None:
             request_timeout_sec=SAFETY_TIMEOUT,
         )
 
+        attrs = {"get_working_batteries.return_value": request.batteries}
+        battery_pool_mock = MagicMock(spec=BatteryPoolStatus, **attrs)
+        BatteryPoolStatus.async_new = AsyncMock(  # type: ignore
+            return_value=battery_pool_mock
+        )
+
         mocker.patch("asyncio.sleep", new_callable=AsyncMock)
         distributor = PowerDistributingActor({"user1": channel.service_handle})
 
-        # Mock that all requested batteries are working.
-        distributor._battery_pool = MagicMock(spec=BatteryPoolStatus)
-        distributor._battery_pool.get_working_batteries.return_value = request.batteries
-
         client_handle = channel.client_handle
         await client_handle.send(request)
 
@@ -175,14 +177,16 @@ async def test_power_distributor_two_users(self, mocker: MockerFixture) -> None:
             "user2": channel2.service_handle,
         }
 
+        attrs = {"get_working_batteries.return_value": {106, 206}}
+        battery_pool_mock = MagicMock(spec=BatteryPoolStatus, **attrs)
+        BatteryPoolStatus.async_new = AsyncMock(  # type: ignore
+            return_value=battery_pool_mock
+        )
+
         mocker.patch("asyncio.sleep", new_callable=AsyncMock)
 
         distributor = PowerDistributingActor(service_channels)
 
-        # Mock that all requested batteries are working.
-        distributor._battery_pool = MagicMock(spec=BatteryPoolStatus)
-        distributor._battery_pool.get_working_batteries.return_value = {106, 206}
-
         user1_handle = channel1.client_handle
         task1 = user1_handle.send(
             Request(
@@ -230,14 +234,15 @@ async def test_power_distributor_invalid_battery_id(
             power=1200, batteries={106, 208}, request_timeout_sec=SAFETY_TIMEOUT
         )
 
+        attrs = {"get_working_batteries.return_value": request.batteries}
+        battery_pool_mock = MagicMock(spec=BatteryPoolStatus, **attrs)
+        BatteryPoolStatus.async_new = AsyncMock(  # type: ignore
+            return_value=battery_pool_mock
+        )
         mocker.patch("asyncio.sleep", new_callable=AsyncMock)
 
         distributor = PowerDistributingActor(service_channels)
 
-        # Mock that all requested batteries are working.
-        distributor._battery_pool = MagicMock(spec=BatteryPoolStatus)
-        distributor._battery_pool.get_working_batteries.return_value = request.batteries
-
         user1_handle = channel1.client_handle
         await user1_handle.send(request)
 
@@ -271,17 +276,16 @@ async def test_power_distributor_overlapping_batteries(
         }
 
         mocker.patch("asyncio.sleep", new_callable=AsyncMock)
+        attrs = {
+            "get_working_batteries.side_effect": [{106, 206}, {106, 306}, {106, 206}]
+        }
+        battery_pool_mock = MagicMock(spec=BatteryPoolStatus, **attrs)
+        BatteryPoolStatus.async_new = AsyncMock(  # type: ignore
+            return_value=battery_pool_mock
+        )
 
         distributor = PowerDistributingActor(service_channels)
 
-        # Mock that all requested batteries are working.
-        distributor._battery_pool = MagicMock(spec=BatteryPoolStatus)
-        distributor._battery_pool.get_working_batteries.side_effect = [
-            {106, 206},
-            {106, 306},
-            {106, 206},
-        ]
-
         user1_handle = channel1.client_handle
         task1 = user1_handle.send(
             Request(
@@ -349,14 +353,16 @@ async def test_power_distributor_one_user_adjust_power_consume(
             adjust_power=False,
         )
 
+        attrs = {"get_working_batteries.return_value": request.batteries}
+        battery_pool_mock = MagicMock(spec=BatteryPoolStatus, **attrs)
+        BatteryPoolStatus.async_new = AsyncMock(  # type: ignore
+            return_value=battery_pool_mock
+        )
+
         mocker.patch("asyncio.sleep", new_callable=AsyncMock)
 
         distributor = PowerDistributingActor(service_channels)
 
-        # Mock that all requested batteries are working.
-        distributor._battery_pool = MagicMock(spec=BatteryPoolStatus)
-        distributor._battery_pool.get_working_batteries.return_value = request.batteries
-
         user1_handle = channel1.client_handle
         await user1_handle.send(request)
 
@@ -394,14 +400,16 @@ async def test_power_distributor_one_user_adjust_power_supply(
             adjust_power=False,
         )
 
+        attrs = {"get_working_batteries.return_value": request.batteries}
+        battery_pool_mock = MagicMock(spec=BatteryPoolStatus, **attrs)
+        BatteryPoolStatus.async_new = AsyncMock(  # type: ignore
+            return_value=battery_pool_mock
+        )
+
         mocker.patch("asyncio.sleep", new_callable=AsyncMock)
 
         distributor = PowerDistributingActor(service_channels)
 
-        # Mock that all requested batteries are working.
-        distributor._battery_pool = MagicMock(spec=BatteryPoolStatus)
-        distributor._battery_pool.get_working_batteries.return_value = request.batteries
-
         user1_handle = channel1.client_handle
         await user1_handle.send(request)
 
@@ -439,14 +447,16 @@ async def test_power_distributor_one_user_adjust_power_success(
             adjust_power=False,
         )
 
+        attrs = {"get_working_batteries.return_value": request.batteries}
+        battery_pool_mock = MagicMock(spec=BatteryPoolStatus, **attrs)
+        BatteryPoolStatus.async_new = AsyncMock(  # type: ignore
+            return_value=battery_pool_mock
+        )
+
         mocker.patch("asyncio.sleep", new_callable=AsyncMock)
 
         distributor = PowerDistributingActor(service_channels)
 
-        # Mock that all requested batteries are working.
-        distributor._battery_pool = MagicMock(spec=BatteryPoolStatus)
-        distributor._battery_pool.get_working_batteries.return_value = request.batteries
-
         user1_handle = channel1.client_handle
         await user1_handle.send(request)
 
@@ -476,16 +486,16 @@ async def test_not_all_batteries_are_working(self, mocker: MockerFixture) -> Non
             power=1200, batteries={106, 206}, request_timeout_sec=SAFETY_TIMEOUT
         )
 
+        attrs = {"get_working_batteries.return_value": request.batteries - {106}}
+        battery_pool_mock = MagicMock(spec=BatteryPoolStatus, **attrs)
+        BatteryPoolStatus.async_new = AsyncMock(  # type: ignore
+            return_value=battery_pool_mock
+        )
+
         mocker.patch("asyncio.sleep", new_callable=AsyncMock)
 
         distributor = PowerDistributingActor({"user1": channel.service_handle})
 
-        # Mock that all requested batteries are working.
-        distributor._battery_pool = MagicMock(spec=BatteryPoolStatus)
-        distributor._battery_pool.get_working_batteries.return_value = (
-            request.batteries - {106}
-        )
-
         client_handle = channel.client_handle
         await client_handle.send(request)