Make the _mocks power distributor test fixture component-agnostic

Then it can be used for testing EV chargers and other component
caterogies as well.

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

Author: shsms

tests/actor/power_distributing/test_power_distributing.py

@@ -70,12 +70,13 @@ class _Mocks:
     streamer: MockComponentDataStreamer
     """A mock component data streamer."""
 
-    battery_status_sender: Sender[ComponentPoolStatus]
-    """Sender for sending status of the batteries."""
+    component_status_sender: Sender[ComponentPoolStatus]
+    """Sender for sending status of the components being tested."""
 
     @classmethod
     async def new(
         cls,
+        component_category: ComponentCategory,
         mocker: MockerFixture,
         graph: _MicrogridComponentGraph | None = None,
         grid_meter: bool | None = None,
@@ -97,11 +98,13 @@ async def new(
         dp = microgrid._data_pipeline._DATA_PIPELINE
         assert dp is not None
 
-        return cls(
-            mockgrid,
-            streamer,
-            dp._battery_power_wrapper.status_channel.new_sender(),
-        )
+        if component_category == ComponentCategory.BATTERY:
+            return cls(
+                mockgrid,
+                streamer,
+                dp._battery_power_wrapper.status_channel.new_sender(),
+            )
+        raise ValueError(f"Unsupported component category: {component_category}")
 
     async def stop(self) -> None:
         """Stop the mocks."""
@@ -120,12 +123,15 @@ async def stop(self) -> None:
 @asynccontextmanager
 async def _mocks(
     mocker: MockerFixture,
+    component_category: ComponentCategory,
     *,
     graph: _MicrogridComponentGraph | None = None,
     grid_meter: bool | None = None,
 ) -> AsyncIterator[_Mocks]:
     """Initialize the mocks."""
-    mocks = await _Mocks.new(mocker, graph=graph, grid_meter=grid_meter)
+    mocks = await _Mocks.new(
+        component_category, mocker, graph=graph, grid_meter=grid_meter
+    )
     try:
         yield mocks
     finally:
@@ -166,7 +172,7 @@ async def _patch_battery_pool_status(
                 ".ComponentPoolStatusTracker",
                 return_value=mock,
             )
-        await mocks.battery_status_sender.send(
+        await mocks.component_status_sender.send(
             ComponentPoolStatus(
                 working=set(mocks.microgrid.battery_ids), uncertain=set()
             )
@@ -262,7 +268,7 @@ async def init_component_data(
 
     async def test_power_distributor_one_user(self, mocker: MockerFixture) -> None:
         """Test if power distribution works with a single user."""
-        mocks = await _Mocks.new(mocker)
+        mocks = await _Mocks.new(ComponentCategory.BATTERY, mocker)
         requests_channel = Broadcast[Request](name="power_distributor requests")
         results_channel = Broadcast[Result](name="power_distributor results")
 
@@ -305,7 +311,7 @@ async def test_power_distributor_exclusion_bounds(
         self, mocker: MockerFixture
     ) -> None:
         """Test if power distributing actor rejects non-zero requests in exclusion bounds."""
-        async with _mocks(mocker) as mocks:
+        async with _mocks(mocker, ComponentCategory.BATTERY) as mocks:
             await self._patch_battery_pool_status(mocks, mocker, {9, 19})
             await self.init_component_data(mocks, skip_batteries={9, 19})
 
@@ -416,7 +422,7 @@ async def test_two_batteries_one_inverters(self, mocker: MockerFixture) -> None:
             )
         )
 
-        async with _mocks(mocker, graph=graph) as mocks:
+        async with _mocks(mocker, ComponentCategory.BATTERY, graph=graph) as mocks:
             await self.init_component_data(mocks)
 
             requests_channel = Broadcast[Request](name="power_distributor requests")
@@ -485,7 +491,7 @@ async def test_two_batteries_one_broken_one_inverters(
             )
         )
 
-        async with _mocks(mocker, graph=graph) as mocks:
+        async with _mocks(mocker, ComponentCategory.BATTERY, graph=graph) as mocks:
             await self.init_component_data(
                 mocks, skip_batteries={bat_components[0].component_id}
             )
@@ -551,7 +557,7 @@ async def test_battery_two_inverters(self, mocker: MockerFixture) -> None:
             )
         )
 
-        async with _mocks(mocker, graph=graph) as mocks:
+        async with _mocks(mocker, ComponentCategory.BATTERY, graph=graph) as mocks:
             await self.init_component_data(mocks)
 
             requests_channel = Broadcast[Request](name="power_distributor requests")
@@ -601,7 +607,7 @@ async def test_two_batteries_three_inverters(self, mocker: MockerFixture) -> Non
             (ComponentCategory.METER, gen.batteries_with_inverter(batteries, 3))
         )
 
-        async with _mocks(mocker, graph=graph) as mocks:
+        async with _mocks(mocker, ComponentCategory.BATTERY, graph=graph) as mocks:
             await self.init_component_data(mocks)
 
             requests_channel = Broadcast[Request](name="power_distributor requests")
@@ -661,7 +667,7 @@ async def test_two_batteries_one_inverter_different_exclusion_bounds_2(
             )
         )
 
-        async with _mocks(mocker, graph=graph) as mocks:
+        async with _mocks(mocker, ComponentCategory.BATTERY, graph=graph) as mocks:
             mocks.streamer.start_streaming(
                 inverter_msg(
                     inverter.component_id,
@@ -750,7 +756,7 @@ async def test_two_batteries_one_inverter_different_exclusion_bounds(
             )
         )
 
-        async with _mocks(mocker, graph=graph) as mocks:
+        async with _mocks(mocker, ComponentCategory.BATTERY, graph=graph) as mocks:
             await self.init_component_data(
                 mocks, skip_batteries={bat.component_id for bat in batteries}
             )
@@ -834,7 +840,7 @@ async def test_connected_but_not_requested_batteries(
             )
         )
 
-        async with _mocks(mocker, graph=graph) as mocks:
+        async with _mocks(mocker, ComponentCategory.BATTERY, graph=graph) as mocks:
             await self.init_component_data(mocks)
 
             requests_channel = Broadcast[Request](name="power_distributor requests")
@@ -881,7 +887,7 @@ async def test_connected_but_not_requested_batteries(
 
     async def test_battery_soc_nan(self, mocker: MockerFixture) -> None:
         """Test if battery with SoC==NaN is not used."""
-        async with _mocks(mocker, grid_meter=False) as mocks:
+        async with _mocks(mocker, ComponentCategory.BATTERY, grid_meter=False) as mocks:
             await self.init_component_data(mocks, skip_batteries={9})
 
             mocks.streamer.start_streaming(
@@ -933,7 +939,7 @@ async def test_battery_soc_nan(self, mocker: MockerFixture) -> None:
 
     async def test_battery_capacity_nan(self, mocker: MockerFixture) -> None:
         """Test battery with capacity set to NaN is not used."""
-        async with _mocks(mocker, grid_meter=False) as mocks:
+        async with _mocks(mocker, ComponentCategory.BATTERY, grid_meter=False) as mocks:
             await self.init_component_data(mocks, skip_batteries={9})
 
             mocks.streamer.start_streaming(
@@ -986,7 +992,7 @@ async def test_battery_capacity_nan(self, mocker: MockerFixture) -> None:
 
     async def test_battery_power_bounds_nan(self, mocker: MockerFixture) -> None:
         """Test battery with power bounds set to NaN is not used."""
-        async with _mocks(mocker, grid_meter=False) as mocks:
+        async with _mocks(mocker, ComponentCategory.BATTERY, grid_meter=False) as mocks:
             await self.init_component_data(
                 mocks, skip_batteries={9}, skip_inverters={8, 18}
             )
@@ -1060,7 +1066,7 @@ async def test_power_distributor_invalid_battery_id(
         self, mocker: MockerFixture
     ) -> None:
         """Test if power distribution raises error if any battery id is invalid."""
-        async with _mocks(mocker, grid_meter=False) as mocks:
+        async with _mocks(mocker, ComponentCategory.BATTERY, grid_meter=False) as mocks:
             await self.init_component_data(mocks)
 
             requests_channel = Broadcast[Request](name="power_distributor requests")
@@ -1101,7 +1107,7 @@ async def test_power_distributor_one_user_adjust_power_consume(
         self, mocker: MockerFixture
     ) -> None:
         """Test if power distribution works with single user works."""
-        async with _mocks(mocker, grid_meter=False) as mocks:
+        async with _mocks(mocker, ComponentCategory.BATTERY, grid_meter=False) as mocks:
             await self.init_component_data(mocks)
 
             requests_channel = Broadcast[Request](name="power_distributor")
@@ -1146,7 +1152,7 @@ async def test_power_distributor_one_user_adjust_power_supply(
         self, mocker: MockerFixture
     ) -> None:
         """Test if power distribution works with single user works."""
-        async with _mocks(mocker, grid_meter=False) as mocks:
+        async with _mocks(mocker, ComponentCategory.BATTERY, grid_meter=False) as mocks:
             await self.init_component_data(mocks)
 
             requests_channel = Broadcast[Request](name="power_distributor requests")
@@ -1191,7 +1197,7 @@ async def test_power_distributor_one_user_adjust_power_success(
         self, mocker: MockerFixture
     ) -> None:
         """Test if power distribution works with single user works."""
-        async with _mocks(mocker, grid_meter=False) as mocks:
+        async with _mocks(mocker, ComponentCategory.BATTERY, grid_meter=False) as mocks:
             await self.init_component_data(mocks)
 
             requests_channel = Broadcast[Request](name="power_distributor requests")
@@ -1234,7 +1240,7 @@ async def test_power_distributor_one_user_adjust_power_success(
 
     async def test_not_all_batteries_are_working(self, mocker: MockerFixture) -> None:
         """Test if power distribution works if not all batteries are working."""
-        async with _mocks(mocker, grid_meter=False) as mocks:
+        async with _mocks(mocker, ComponentCategory.BATTERY, grid_meter=False) as mocks:
             await self.init_component_data(mocks)
 
             batteries = {9, 19}
@@ -1278,7 +1284,7 @@ async def test_not_all_batteries_are_working(self, mocker: MockerFixture) -> Non
 
     async def test_partial_failure_result(self, mocker: MockerFixture) -> None:
         """Test power results when the microgrid failed to set power for one of the batteries."""
-        async with _mocks(mocker, grid_meter=False) as mocks:
+        async with _mocks(mocker, ComponentCategory.BATTERY, grid_meter=False) as mocks:
             await self.init_component_data(mocks)
 
             batteries = {9, 19, 29}