Update BatteryPool tests, examples to use microgrid.battery_pool() (#336)

... instead of creating their own `BatteryPool` instances.

Author: sahas-subramanian-frequenz

examples/battery_pool.py

@@ -7,71 +7,18 @@
 
 import asyncio
 import logging
-from datetime import timedelta
 from typing import Any, Dict
 
-from frequenz.channels import Broadcast, Receiver
+from frequenz.channels import Receiver
 from frequenz.channels.util import MergeNamed
 
 from frequenz.sdk import microgrid
-from frequenz.sdk.actor import (
-    ChannelRegistry,
-    ComponentMetricRequest,
-    DataSourcingActor,
-    ResamplerConfig,
-)
-from frequenz.sdk.actor.power_distributing import PowerDistributingActor
-from frequenz.sdk.actor.power_distributing._battery_pool_status import BatteryStatus
-from frequenz.sdk.microgrid import connection_manager
-from frequenz.sdk.microgrid.component import ComponentCategory
-from frequenz.sdk.timeseries.battery_pool import BatteryPool
+from frequenz.sdk.actor import ResamplerConfig
 
 HOST = "microgrid.sandbox.api.frequenz.io"  # it should be the host name.
 PORT = 61060
 
 
-def create_battery_pool() -> BatteryPool:
-    """Create battery pool.
-
-    It needs many instance to be created before.
-
-    Returns:
-        BatteryPool instance ready to use.
-    """
-    channel_registry = ChannelRegistry(name="data-registry")
-
-    # Create a channels for sending/receiving subscription requests
-    data_source_request_channel = Broadcast[ComponentMetricRequest](
-        "data-source", resend_latest=True
-    )
-
-    # Instantiate a data sourcing actor
-    _ = DataSourcingActor(
-        request_receiver=data_source_request_channel.new_receiver(
-            "data_sourcing_receiver"
-        ),
-        registry=channel_registry,
-    )
-
-    battery_status_channel = Broadcast[BatteryStatus]("batteries-status")
-    _ = PowerDistributingActor(
-        users_channels={},
-        battery_status_sender=battery_status_channel.new_sender(),
-    )
-
-    batteries = connection_manager.get().component_graph.components(
-        component_category={ComponentCategory.BATTERY}
-    )
-
-    return BatteryPool(
-        batteries_id=set(battery.component_id for battery in batteries),
-        batteries_status_receiver=battery_status_channel.new_receiver(
-            name="battery_pool", maxsize=1
-        ),
-        min_update_interval=timedelta(seconds=0.2),
-    )
-
-
 async def main() -> None:
     """Create the battery pool, activate all formulas and listen for any update."""
     logging.basicConfig(
@@ -81,7 +28,7 @@ async def main() -> None:
         host=HOST, port=PORT, resampler_config=ResamplerConfig(resampling_period_s=1.0)
     )
 
-    battery_pool = create_battery_pool()
+    battery_pool = microgrid.battery_pool()
     receivers: Dict[str, Receiver[Any]] = {
         "soc": await battery_pool.soc(maxsize=1),
         "capacity": await battery_pool.capacity(maxsize=1),

src/frequenz/sdk/microgrid/__init__.py

@@ -9,7 +9,7 @@
 
 from ..actor import ResamplerConfig
 from . import _data_pipeline, client, component, connection_manager
-from ._data_pipeline import ev_charger_pool, logical_meter
+from ._data_pipeline import battery_pool, ev_charger_pool, logical_meter
 from ._graph import ComponentGraph
 
 
@@ -30,6 +30,7 @@ async def initialize(host: str, port: int, resampler_config: ResamplerConfig) ->
     "initialize",
     "client",
     "component",
+    "battery_pool",
     "ev_charger_pool",
     "logical_meter",
 ]

src/frequenz/sdk/microgrid/_data_pipeline.py

@@ -81,7 +81,9 @@ def __init__(
         self._data_sourcing_actor: _ActorInfo | None = None
         self._resampling_actor: _ActorInfo | None = None
 
-        self._battery_status_channel = Broadcast["BatteryStatus"]("battery-status")
+        self._battery_status_channel = Broadcast["BatteryStatus"](
+            "battery-status", resend_latest=True
+        )
         self._power_distribution_channel = Bidirectional["Request", "Result"](
             "Default", "Power Distributing Actor"
         )
@@ -170,7 +172,9 @@ def battery_pool(
 
         if key not in self._battery_pools:
             self._battery_pools[key] = BatteryPool(
-                batteries_status_receiver=self._battery_status_channel.new_receiver(),
+                batteries_status_receiver=self._battery_status_channel.new_receiver(
+                    maxsize=1
+                ),
                 min_update_interval=timedelta(
                     seconds=self._resampler_config.resampling_period_s
                 ),

tests/timeseries/_battery_pool/test_battery_pool.py

@@ -8,18 +8,20 @@
 import dataclasses
 import logging
 from dataclasses import dataclass, is_dataclass, replace
-from datetime import datetime, timedelta, timezone
+from datetime import datetime, timezone
 from typing import Any, AsyncIterator, Generic, Iterator, TypeVar
 
 import async_solipsism
 import pytest
-from frequenz.channels import Broadcast, Receiver, Sender
+from frequenz.channels import Receiver, Sender
 from pytest_mock import MockerFixture
 
+from frequenz.sdk import microgrid
 from frequenz.sdk._internal._constants import (
     MAX_BATTERY_DATA_AGE_SEC,
     WAIT_FOR_COMPONENT_DATA_SEC,
 )
+from frequenz.sdk.actor import ResamplerConfig
 from frequenz.sdk.actor.power_distributing import BatteryStatus
 from frequenz.sdk.microgrid.component import ComponentCategory
 from frequenz.sdk.timeseries.battery_pool import (
@@ -40,7 +42,6 @@
     create_component_graph_structure,
 )
 from ...utils.mock_microgrid_client import MockMicrogridClient
-from ...utils.sdk_interface import SdkInterface
 
 _logger = logging.getLogger(__name__)
 
@@ -130,27 +131,37 @@ async def setup_all_batteries(mocker: MockerFixture) -> AsyncIterator[SetupArgs]
     mock_microgrid = create_mock_microgrid(
         mocker, ComponentGraphConfig(batteries_num=2)
     )
-    sdk = SdkInterface(resampling_period_s=1)
+    min_update_interval: float = 0.2
+    # pylint: disable=protected-access
+    microgrid._data_pipeline._DATA_PIPELINE = None
+    microgrid._data_pipeline.initialize(ResamplerConfig(min_update_interval))
     streamer = MockComponentDataStreamer(mock_microgrid)
 
     # We don't use status channel from the sdk interface to limit
     # the scope of this tests. This tests should cover BatteryPool only.
     # We use our own battery status channel, where we easily control set of working
     # batteries.
-    battery_status_channel = Broadcast[BatteryStatus]("bat_status", resend_latest=True)
-    sender_channel = battery_status_channel.new_sender()
-    min_update_interval: float = 0.2
-    battery_pool = BatteryPool(
-        batteries_status_receiver=battery_status_channel.new_receiver(maxsize=1),
-        min_update_interval=timedelta(seconds=min_update_interval),
-    )
+    battery_pool = microgrid.battery_pool()
+
+    assert microgrid._data_pipeline._DATA_PIPELINE is not None
 
     args = SetupArgs(
-        battery_pool, min_update_interval, mock_microgrid, streamer, sender_channel
+        battery_pool,
+        min_update_interval,
+        mock_microgrid,
+        streamer,
+        microgrid._data_pipeline._DATA_PIPELINE._battery_status_channel.new_sender(),
     )
 
     yield args
-    await asyncio.gather(*[sdk.stop(), battery_pool.stop(), streamer.stop()])
+    await asyncio.gather(
+        *[
+            microgrid._data_pipeline._DATA_PIPELINE._stop(),
+            battery_pool.stop(),
+            streamer.stop(),
+        ]
+    )
+    # pylint: enable=protected-access
 
 
 @pytest.fixture
@@ -170,29 +181,39 @@ async def setup_batteries_pool(mocker: MockerFixture) -> AsyncIterator[SetupArgs
         mocker, ComponentGraphConfig(batteries_num=4)
     )
     streamer = MockComponentDataStreamer(mock_microgrid)
-    sdk = SdkInterface(resampling_period_s=1)
+    min_update_interval: float = 0.2
+    # pylint: disable=protected-access
+    microgrid._data_pipeline._DATA_PIPELINE = None
+    microgrid._data_pipeline.initialize(ResamplerConfig(min_update_interval))
 
     # We don't use status channel from the sdk interface to limit
     # the scope of this tests. This tests should cover BatteryPool only.
     # We use our own battery status channel, where we easily control set of working
     # batteries.
-    battery_status_channel = Broadcast[BatteryStatus]("bat_status", resend_latest=True)
-    sender_channel = battery_status_channel.new_sender()
-    min_update_interval: float = 0.2
     all_batteries = list(get_components(mock_microgrid, ComponentCategory.BATTERY))
 
-    battery_pool = BatteryPool(
-        batteries_id=set(all_batteries[:2]),
-        batteries_status_receiver=battery_status_channel.new_receiver(maxsize=1),
-        min_update_interval=timedelta(seconds=min_update_interval),
-    )
+    battery_pool = microgrid.battery_pool(set(all_batteries[:2]))
+
+    assert microgrid._data_pipeline._DATA_PIPELINE is not None
 
     args = SetupArgs(
-        battery_pool, min_update_interval, mock_microgrid, streamer, sender_channel
+        battery_pool,
+        min_update_interval,
+        mock_microgrid,
+        streamer,
+        microgrid._data_pipeline._DATA_PIPELINE._battery_status_channel.new_sender(),
     )
 
     yield args
-    await asyncio.gather(*[sdk.stop(), battery_pool.stop(), streamer.stop()])
+
+    await asyncio.gather(
+        *[
+            microgrid._data_pipeline._DATA_PIPELINE._stop(),
+            battery_pool.stop(),
+            streamer.stop(),
+        ]
+    )
+    # pylint: enable=protected-access
 
 
 T = TypeVar("T")

tests/timeseries/_formula_engine/test_formula_composition.py

@@ -23,7 +23,7 @@ async def test_formula_composition(  # pylint: disable=too-many-locals
         mocker: MockerFixture,
     ) -> None:
         """Test the composition of formulas."""
-        mockgrid = MockMicrogrid(grid_side_meter=False, sample_rate_s=0.1)
+        mockgrid = MockMicrogrid(grid_side_meter=False, sample_rate_s=0.15)
         mockgrid.add_batteries(3)
         mockgrid.add_solar_inverters(2)
         await mockgrid.start(mocker)