Tests: Add convenience methods to create batteries with multiple inverters

Signed-off-by: Mathias L. Baumann <[email protected]>

Author: Marenz

tests/actor/power_distributing/test_power_distributing.py

@@ -537,21 +537,7 @@ async def test_battery_two_inverters(self, mocker: MockerFixture) -> None:
         graph = gen.to_graph(
             (
                 ComponentCategory.METER,
-                [
-                    (
-                        ComponentCategory.METER,
-                        [
-                            (
-                                ComponentCategory.INVERTER,
-                                bat_component,
-                            ),
-                            (
-                                ComponentCategory.INVERTER,
-                                bat_component,
-                            ),
-                        ],
-                    ),
-                ],
+                [gen.battery_with_inverter(bat_component, 2)],
             )
         )
 
@@ -600,28 +586,7 @@ async def test_two_batteries_three_inverters(self, mocker: MockerFixture) -> Non
         batteries = gen.components(*[ComponentCategory.BATTERY] * 2)
 
         graph = gen.to_graph(
-            (
-                ComponentCategory.METER,
-                [
-                    (
-                        ComponentCategory.METER,
-                        [
-                            (
-                                ComponentCategory.INVERTER,
-                                [*batteries],
-                            ),
-                            (
-                                ComponentCategory.INVERTER,
-                                [*batteries],
-                            ),
-                            (
-                                ComponentCategory.INVERTER,
-                                [*batteries],
-                            ),
-                        ],
-                    ),
-                ],
-            )
+            (ComponentCategory.METER, gen.batteries_with_inverter(batteries, 3))
         )
 
         async with _mocks(mocker, graph=graph) as mocks:

tests/utils/graph_generator.py

@@ -70,6 +70,55 @@ def component(self, other: Component) -> Component:
             the given component.
         """
 
+    def battery_with_inverter(self, battery: Component, num_inverters: int) -> Any:
+        """Add a meter and inverters to the given battery.
+
+        Args:
+            one_or_more_batteries: the battery component.
+            num_inverters: the number of inverters to create.
+
+        Returns:
+            connected graph components for the given battery.
+        """
+        return self._battery_with_inverter(battery, num_inverters)
+
+    def batteries_with_inverter(
+        self, one_or_more_batteries: list[Component], num_inverters: int
+    ) -> Any:
+        """Add a meter and inverters to the given batteries.
+
+        Args:
+            one_or_more_batteries: the battery components.
+            num_inverters: the number of inverters to create.
+
+        Returns:
+            connected graph components for the given batteries.
+        """
+        return self._battery_with_inverter(one_or_more_batteries, num_inverters)
+
+    def _battery_with_inverter(
+        self, one_or_more_batteries: Component | list[Component], num_inverters: int
+    ) -> Any:
+        """Add a meter and inverters to the given battery or batteries.
+
+        Args:
+            one_or_more_batteries: the battery component or components.
+            num_inverters: the number of inverters to create.
+
+        Returns:
+            connected graph components for the given battery or batteries.
+        """
+        return (
+            ComponentCategory.METER,
+            [
+                (
+                    self.component(ComponentCategory.INVERTER, InverterType.BATTERY),
+                    one_or_more_batteries,
+                )
+                for _ in range(num_inverters)
+            ],
+        )
+
     @overload
     def component(
         self, other: ComponentCategory, comp_type: ComponentType | None = None