Allow configuring the component graph from microgrid.initialize()

The new `component_graph_config` keyword argument takes a
`ComponentGraphConfig` and is passed to the `ComponentGraph`
constructor.  When it is `None`, the library's default configuration
is used.  `ComponentGraphConfig` and `FormulaOverrides` are
re-exported from `frequenz.sdk.microgrid`.

Signed-off-by: Sahas Subramanian <sahas.subramanian@proton.me>

Author: shsms

RELEASE_NOTES.md

@@ -12,7 +12,7 @@
 
 ## New Features
 
-<!-- Here goes the main new features and examples or instructions on how to use them -->
+* `microgrid.initialize()` has a new keyword argument `component_graph_config`.  It takes a `ComponentGraphConfig` and gives full control over how the component graph is built.  For example, pass `ComponentGraphConfig(prefer_meters_in_component_formulas=True)` to make the per-category formulas read from the meter first, as in previous releases.  `ComponentGraphConfig` and `FormulaOverrides` are re-exported from `frequenz.sdk.microgrid`.
 
 ## Bug Fixes
 

src/frequenz/sdk/microgrid/__init__.py

@@ -370,6 +370,8 @@
 
 from datetime import timedelta
 
+from frequenz.microgrid_component_graph import ComponentGraphConfig, FormulaOverrides
+
 from ..timeseries._resampling._config import ResamplerConfig
 from . import _data_pipeline, connection_manager
 from ._data_pipeline import (
@@ -393,6 +395,7 @@ async def initialize(
     api_power_request_timeout: timedelta = timedelta(seconds=5.0),
     # pylint: disable-next: line-too-long
     battery_power_manager_algorithm: PowerManagerAlgorithm = PowerManagerAlgorithm.SHIFTING_MATRYOSHKA,  # noqa: E501
+    component_graph_config: ComponentGraphConfig | None = None,
 ) -> None:
     """Initialize the microgrid connection manager and the data pipeline.
 
@@ -409,8 +412,17 @@ async def initialize(
             will be unavailable from the corresponding component pools.
         battery_power_manager_algorithm: The power manager algorithm to use for
             batteries.
+        component_graph_config: The configuration for building the component
+            graph.  Use it, for example, to make the per-category formulas
+            (like the battery or PV formulas) read from the meter first
+            (`prefer_meters_in_component_formulas`), or to set per-formula
+            overrides through `FormulaOverrides`.  When `None`, the library's
+            default configuration is used.
     """
-    await connection_manager.initialize(server_url)
+    await connection_manager.initialize(
+        server_url,
+        component_graph_config=component_graph_config,
+    )
     await _data_pipeline.initialize(
         resampler_config,
         api_power_request_timeout=api_power_request_timeout,
@@ -419,6 +431,8 @@ async def initialize(
 
 
 __all__ = [
+    "ComponentGraphConfig",
+    "FormulaOverrides",
     "PowerManagerAlgorithm",
     "initialize",
     "consumer",

src/frequenz/sdk/microgrid/connection_manager.py

@@ -18,6 +18,7 @@
     MicrogridApiClient,
     MicrogridInfo,
 )
+from frequenz.microgrid_component_graph import ComponentGraphConfig
 
 from .component_graph import ComponentGraph
 
@@ -27,7 +28,12 @@
 class ConnectionManager(ABC):
     """Creates and stores core features."""
 
-    def __init__(self, server_url: str) -> None:
+    def __init__(
+        self,
+        server_url: str,
+        *,
+        component_graph_config: ComponentGraphConfig | None = None,
+    ) -> None:
         """Create object instance.
 
         Args:
@@ -36,9 +42,12 @@ def __init__(self, server_url: str) -> None:
                 where the port should be an int between `0` and `65535` (defaulting to
                 `9090`) and ssl should be a boolean (defaulting to false). For example:
                 `grpc://localhost:1090?ssl=true`.
+            component_graph_config: The configuration for building the component
+                graph.  When `None`, the library's default configuration is used.
         """
         super().__init__()
         self._server_url = server_url
+        self._component_graph_config = component_graph_config
 
     @property
     def server_url(self) -> str:
@@ -94,7 +103,12 @@ async def _initialize(self) -> None:
 class _InsecureConnectionManager(ConnectionManager):
     """Microgrid Api with insecure channel implementation."""
 
-    def __init__(self, server_url: str) -> None:
+    def __init__(
+        self,
+        server_url: str,
+        *,
+        component_graph_config: ComponentGraphConfig | None = None,
+    ) -> None:
         """Create and stores core features.
 
         Args:
@@ -103,8 +117,10 @@ def __init__(self, server_url: str) -> None:
                 where the port should be an int between `0` and `65535` (defaulting to
                 `9090`) and ssl should be a boolean (defaulting to false). For example:
                 `grpc://localhost:1090?ssl=true`.
+            component_graph_config: The configuration for building the component
+                graph.  When `None`, the library's default configuration is used.
         """
-        super().__init__(server_url)
+        super().__init__(server_url, component_graph_config=component_graph_config)
         self._client = MicrogridApiClient(server_url)
         # To create graph from the API client we need await.
         # So create empty graph here, and update it in `run` method.
@@ -172,14 +188,23 @@ async def _initialize(self) -> None:
             self._client.list_components(),
             self._client.list_connections(),
         )
-        self._graph = ComponentGraph(set(components), set(connections))
+        if self._component_graph_config is None:
+            self._graph = ComponentGraph(set(components), set(connections))
+        else:
+            self._graph = ComponentGraph(
+                set(components), set(connections), self._component_graph_config
+            )
 
 
 _CONNECTION_MANAGER: ConnectionManager | None = None
 """The ConnectionManager singleton instance."""
 
 
-async def initialize(server_url: str) -> None:
+async def initialize(
+    server_url: str,
+    *,
+    component_graph_config: ComponentGraphConfig | None = None,
+) -> None:
     """Initialize the MicrogridApi. This function should be called only once.
 
     Args:
@@ -188,6 +213,8 @@ async def initialize(server_url: str) -> None:
             where the port should be an int between `0` and `65535` (defaulting to
             `9090`) and ssl should be a boolean (defaulting to false). For example:
             `grpc://localhost:1090?ssl=true`.
+        component_graph_config: The configuration for building the component
+            graph.  When `None`, the library's default configuration is used.
     """
     # From Doc: pylint just try to discourage this usage.
     # That doesn't mean you cannot use it.
@@ -197,7 +224,10 @@ async def initialize(server_url: str) -> None:
 
     _logger.info("Connecting to microgrid at %s", server_url)
 
-    connection_manager = _InsecureConnectionManager(server_url)
+    connection_manager = _InsecureConnectionManager(
+        server_url,
+        component_graph_config=component_graph_config,
+    )
     await connection_manager._initialize()  # pylint: disable=protected-access
 
     # Check again that _MICROGRID_API is None in case somebody had the great idea of

tests/microgrid/test_microgrid_api.py

@@ -28,7 +28,7 @@
     Meter,
 )
 
-from frequenz.sdk.microgrid import connection_manager
+from frequenz.sdk.microgrid import ComponentGraphConfig, connection_manager
 
 _MICROGRID_ID = MicrogridId(1)
 
@@ -135,6 +135,28 @@ def microgrid(self) -> MicrogridInfo:
             ),
         )
 
+    @staticmethod
+    def _mock_microgrid_client(
+        components: list[list[Component]],
+        connections: list[list[ComponentConnection]],
+        microgrid: MicrogridInfo,
+    ) -> MagicMock:
+        """Create a mock microgrid API client.
+
+        Args:
+            components: components to return, one list per call.
+            connections: connections to return, one list per call.
+            microgrid: the information about the microgrid.
+
+        Returns:
+            the mock client.
+        """
+        client = MagicMock()
+        client.list_components = AsyncMock(side_effect=components)
+        client.list_connections = AsyncMock(side_effect=connections)
+        client.get_microgrid_info = AsyncMock(return_value=microgrid)
+        return client
+
     @mock.patch("grpc.aio.insecure_channel")
     async def test_connection_manager(
         self,
@@ -151,10 +173,9 @@ async def test_connection_manager(
             connections: connections
             microgrid: the information about the microgrid
         """
-        microgrid_client = MagicMock()
-        microgrid_client.list_components = AsyncMock(side_effect=components)
-        microgrid_client.list_connections = AsyncMock(side_effect=connections)
-        microgrid_client.get_microgrid_info = AsyncMock(return_value=microgrid)
+        microgrid_client = self._mock_microgrid_client(
+            components, connections, microgrid
+        )
 
         with mock.patch(
             "frequenz.sdk.microgrid.connection_manager.MicrogridApiClient",
@@ -214,6 +235,57 @@ async def test_connection_manager(
             assert api.microgrid_id == microgrid.id
             assert api.location == microgrid.location
 
+    @mock.patch("grpc.aio.insecure_channel")
+    async def test_component_graph_config(
+        self,
+        _insecure_channel_mock: MagicMock,
+        components: list[list[Component]],
+        connections: list[list[ComponentConnection]],
+        microgrid: MicrogridInfo,
+    ) -> None:
+        """Test that the component graph config is used to build the graph.
+
+        Args:
+            _insecure_channel_mock: insecure channel mock from `mock.patch`
+            components: components
+            connections: connections
+            microgrid: the information about the microgrid
+        """
+        microgrid_client = self._mock_microgrid_client(
+            [components[0]] * 2, [connections[0]] * 2, microgrid
+        )
+
+        with mock.patch(
+            "frequenz.sdk.microgrid.connection_manager.MicrogridApiClient",
+            return_value=microgrid_client,
+        ):
+            # pylint: disable=protected-access
+            default_manager = connection_manager._InsecureConnectionManager(
+                "grpc://127.0.0.1:10001"
+            )
+            await default_manager._initialize()
+
+            meters_first_manager = connection_manager._InsecureConnectionManager(
+                "grpc://127.0.0.1:10001",
+                component_graph_config=ComponentGraphConfig(
+                    prefer_meters_in_component_formulas=True
+                ),
+            )
+            await meters_first_manager._initialize()
+
+        batteries = {ComponentId(9), ComponentId(12)}
+        # By default the battery inverters are the primary source and the
+        # meters are the fallback.
+        assert (
+            default_manager.component_graph.battery_formula(batteries)
+            == "COALESCE(#8, #7, 0.0) + COALESCE(#11, #10, 0.0)"
+        )
+        # With `prefer_meters_in_component_formulas` the meters come first.
+        assert (
+            meters_first_manager.component_graph.battery_formula(batteries)
+            == "COALESCE(#7, #8, 0.0) + COALESCE(#10, #11, 0.0)"
+        )
+
     @mock.patch("grpc.aio.insecure_channel")
     async def test_connection_manager_another_method(
         self,