Return BatteryPoolWrapper from the data pipeline instead of BatteryPool

This allows users to get a custom object that will keep track of the
user's source_id and priorities, without having to duplicate the
formula engine and metric calculators.

Signed-off-by: Sahas Subramanian <[email protected]>

Author: shsms

benchmarks/power_distribution/power_distributor.py

@@ -51,7 +51,7 @@ async def send_requests(batteries: set[int], request_num: int) -> list[Result]:
         List of the results from the PowerDistributingActor.
     """
     battery_pool = microgrid.battery_pool(batteries)
-    results_rx = battery_pool.power_bounds().new_receiver()
+    results_rx = battery_pool.power_bounds.new_receiver()
     result: list[Any] = []
     for _ in range(request_num):
         await battery_pool.set_power(Power(float(random.randrange(100000, 1000000))))

src/frequenz/sdk/microgrid/_data_pipeline.py

@@ -37,7 +37,7 @@
         Request,
         Result,
     )
-    from ..timeseries.battery_pool import BatteryPool
+    from ..timeseries.battery_pool import BatteryPool, BatteryPoolWrapper
     from ..timeseries.ev_charger_pool import EVChargerPool
     from ..timeseries.logical_meter import LogicalMeter
 
@@ -185,20 +185,26 @@ def ev_charger_pool(
     def battery_pool(
         self,
         battery_ids: abc.Set[int] | None = None,
-    ) -> BatteryPool:
+        source_id: str | None = None,
+        priority: int = 0,
+    ) -> BatteryPoolWrapper:
         """Return the corresponding BatteryPool instance for the given ids.
 
         If a BatteryPool instance for the given ids doesn't exist, a new one is created
         and returned.
 
+        The BatteryPool is wrapped in a new `BatteryPoolWrapper` instance each time.
+
         Args:
             battery_ids: Optional set of IDs of batteries to be managed by the
                 BatteryPool.
+            source_id: The source ID to use for the requests made with this instance.
+            priority: The priority of the actor making the call.
 
         Returns:
-            A BatteryPool instance.
+            A BatteryPoolWrapper instance.
         """
-        from ..timeseries.battery_pool import BatteryPool
+        from ..timeseries.battery_pool import BatteryPool, BatteryPoolWrapper
 
         if not self._power_managing_actor:
             self._start_power_managing_actor()
@@ -225,7 +231,7 @@ def battery_pool(
                 batteries_id=battery_ids,
             )
 
-        return self._battery_pools[key]
+        return BatteryPoolWrapper(self._battery_pools[key], source_id, priority)
 
     def _start_power_managing_actor(self) -> None:
         """Start the power managing actor if it is not already running."""
@@ -415,21 +421,29 @@ def ev_charger_pool(ev_charger_ids: set[int] | None = None) -> EVChargerPool:
     return _get().ev_charger_pool(ev_charger_ids)
 
 
-def battery_pool(battery_ids: abc.Set[int] | None = None) -> BatteryPool:
+def battery_pool(
+    battery_ids: abc.Set[int] | None = None,
+    source_id: str | None = None,
+    priority: int = 0,
+) -> BatteryPoolWrapper:
     """Return the corresponding BatteryPool instance for the given ids.
 
     If a BatteryPool instance for the given ids doesn't exist, a new one is
     created and returned.
 
+    The BatteryPool is wrapped in a new `BatteryPoolWrapper` instance each time.
+
     Args:
         battery_ids: Optional set of IDs of batteries to be managed by the
             BatteryPool.  If not specified, all batteries available in the
             component graph are used.
+        source_id: The source ID to use for the requests made with this instance.
+        priority: The priority of the actor making the call.
 
     Returns:
         A BatteryPool instance.
     """
-    return _get().battery_pool(battery_ids)
+    return _get().battery_pool(battery_ids, source_id, priority)
 
 
 def _get() -> _DataPipeline:

src/frequenz/sdk/timeseries/battery_pool/__init__.py

@@ -3,10 +3,12 @@
 
 """Manage a pool of batteries."""
 
+from ._battery_pool_wrapper import BatteryPoolWrapper
 from ._result_types import PowerMetrics
 from .battery_pool import BatteryPool
 
 __all__ = [
     "BatteryPool",
+    "BatteryPoolWrapper",
     "PowerMetrics",
 ]

tests/timeseries/_battery_pool/test_battery_pool.py

@@ -36,7 +36,7 @@
     Sample,
     Temperature,
 )
-from frequenz.sdk.timeseries.battery_pool import BatteryPool, PowerMetrics
+from frequenz.sdk.timeseries.battery_pool import BatteryPoolWrapper, PowerMetrics
 from frequenz.sdk.timeseries.battery_pool._metric_calculator import (
     battery_inverter_mapping,
 )
@@ -90,7 +90,7 @@ def get_components(
 class SetupArgs:
     """Setup arguments needed to run tests."""
 
-    battery_pool: BatteryPool
+    battery_pool: BatteryPoolWrapper
     """Battery pool that should be tested."""
 
     min_update_interval: float
@@ -168,7 +168,7 @@ async def setup_all_batteries(mocker: MockerFixture) -> AsyncIterator[SetupArgs]
     await asyncio.gather(
         *[
             microgrid._data_pipeline._DATA_PIPELINE._stop(),
-            battery_pool.stop(),
+            battery_pool._battery_pool.stop(),
             streamer.stop(),
         ]
     )
@@ -222,7 +222,7 @@ async def setup_batteries_pool(mocker: MockerFixture) -> AsyncIterator[SetupArgs
     await asyncio.gather(
         *[
             microgrid._data_pipeline._DATA_PIPELINE._stop(),
-            battery_pool.stop(),
+            battery_pool._battery_pool.stop(),
             streamer.stop(),
         ]
     )

tests/timeseries/_battery_pool/test_battery_pool_control_methods.py

@@ -185,7 +185,7 @@ async def test_case_1(
         #
         # It will be replaced by a reporting streaming from the PowerManager in a
         # subsequent commit.
-        bounds_rx = battery_pool.power_bounds().new_receiver()
+        bounds_rx = battery_pool.power_bounds.new_receiver()
 
         self._assert_report(
             await bounds_rx.receive(), power=None, lower=-4000.0, upper=4000.0
@@ -279,9 +279,9 @@ async def test_case_2(self, mocks: Mocks, mocker: MockerFixture) -> None:
         await self._init_data_for_inverters(mocks)
 
         battery_pool_1 = microgrid.battery_pool(set(mocks.microgrid.battery_ids[:2]))
-        bounds_1_rx = battery_pool_1.power_bounds().new_receiver()
+        bounds_1_rx = battery_pool_1.power_bounds.new_receiver()
         battery_pool_2 = microgrid.battery_pool(set(mocks.microgrid.battery_ids[2:]))
-        bounds_2_rx = battery_pool_2.power_bounds().new_receiver()
+        bounds_2_rx = battery_pool_2.power_bounds.new_receiver()
 
         self._assert_report(
             await bounds_1_rx.receive(), power=None, lower=-2000.0, upper=2000.0
@@ -324,10 +324,10 @@ async def test_case_3(self, mocks: Mocks, mocker: MockerFixture) -> None:
         await self._init_data_for_batteries(mocks)
         await self._init_data_for_inverters(mocks)
 
-        battery_pool_1 = microgrid.battery_pool()
-        bounds_1_rx = battery_pool_1.power_bounds(2).new_receiver()
-        battery_pool_2 = microgrid.battery_pool()
-        bounds_2_rx = battery_pool_2.power_bounds(1).new_receiver()
+        battery_pool_1 = microgrid.battery_pool(priority=2)
+        bounds_1_rx = battery_pool_1.power_bounds.new_receiver()
+        battery_pool_2 = microgrid.battery_pool(priority=1)
+        bounds_2_rx = battery_pool_2.power_bounds.new_receiver()
 
         self._assert_report(
             await bounds_1_rx.receive(), power=None, lower=-4000.0, upper=4000.0
@@ -337,7 +337,6 @@ async def test_case_3(self, mocks: Mocks, mocker: MockerFixture) -> None:
         )
         await battery_pool_1.set_power(
             Power.from_watts(-1000.0),
-            _priority=2,
             _bounds=(Power.from_watts(-1000.0), Power.from_watts(0.0)),
         )
         self._assert_report(
@@ -356,7 +355,6 @@ async def test_case_3(self, mocks: Mocks, mocker: MockerFixture) -> None:
 
         await battery_pool_2.set_power(
             Power.from_watts(0.0),
-            _priority=1,
             _bounds=(Power.from_watts(0.0), Power.from_watts(1000.0)),
         )
         self._assert_report(

tests/timeseries/_formula_engine/test_formula_composition.py

@@ -94,7 +94,7 @@ async def test_formula_composition(  # pylint: disable=too-many-locals
 
         await mockgrid.cleanup()
         await engine._stop()  # pylint: disable=protected-access
-        await battery_pool.stop()
+        await battery_pool._battery_pool.stop()  # pylint: disable=protected-access
         await logical_meter.stop()
 
     async def test_formula_composition_missing_pv(self, mocker: MockerFixture) -> None:
@@ -130,7 +130,7 @@ async def test_formula_composition_missing_pv(self, mocker: MockerFixture) -> No
 
         await mockgrid.cleanup()
         await engine._stop()  # pylint: disable=protected-access
-        await battery_pool.stop()
+        await battery_pool._battery_pool.stop()  # pylint: disable=protected-access
         await logical_meter.stop()
 
         assert count == 10
@@ -165,7 +165,7 @@ async def test_formula_composition_missing_bat(self, mocker: MockerFixture) -> N
 
         await mockgrid.cleanup()
         await engine._stop()  # pylint: disable=protected-access
-        await battery_pool.stop()
+        await battery_pool._battery_pool.stop()  # pylint: disable=protected-access
         await logical_meter.stop()
 
         assert count == 10