Add tests for the ComponentGraphGenerator

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

Author: shsms

tests/test_gridpool.py

@@ -3,9 +3,66 @@
 
 """Tests for the frequenz.gridpool package."""
 
-import frequenz.gridpool
+import asyncio
+from unittest.mock import AsyncMock, MagicMock, patch
 
+from frequenz.client.assets import AssetsApiClient
+from frequenz.client.assets.electrical_component import (
+    ComponentConnection,
+    GridConnectionPoint,
+    Meter,
+    SolarInverter,
+)
+from frequenz.client.common.microgrid import MicrogridId
+from frequenz.client.common.microgrid.electrical_components import ElectricalComponentId
 
-def test_loading() -> None:
+from frequenz.gridpool._graph_generator import ComponentGraphGenerator
+
+
+async def test_formula_generation() -> None:
     """Test that the frequenz.gridpool package loads correctly."""
-    assert frequenz.gridpool is not None
+    assets_client_mock = MagicMock(spec=AssetsApiClient)
+    assets_client_mock.list_microgrid_electrical_components = AsyncMock(
+        return_value=[
+            GridConnectionPoint(
+                id=ElectricalComponentId(1),
+                microgrid_id=MicrogridId(10),
+                rated_fuse_current=100,
+            ),
+            Meter(
+                id=ElectricalComponentId(2),
+                microgrid_id=MicrogridId(10),
+            ),
+            Meter(
+                id=ElectricalComponentId(3),
+                microgrid_id=MicrogridId(10),
+            ),
+            SolarInverter(
+                id=ElectricalComponentId(4),
+                microgrid_id=MicrogridId(10),
+            ),
+        ]
+    )
+    assets_client_mock.list_microgrid_electrical_component_connections = AsyncMock(
+        return_value=[
+            ComponentConnection(
+                source=ElectricalComponentId(1),
+                destination=ElectricalComponentId(2),
+            ),
+            ComponentConnection(
+                source=ElectricalComponentId(1),
+                destination=ElectricalComponentId(3),
+            ),
+            ComponentConnection(
+                source=ElectricalComponentId(2),
+                destination=ElectricalComponentId(4),
+            ),
+        ]
+    )
+
+    g = ComponentGraphGenerator("grpc://never.where:-55", connect=False)
+    with patch.object(g, "_client", assets_client_mock):
+        graph = await g.get_component_graph(MicrogridId(10))
+
+        assert graph.grid_formula() == "COALESCE(#2, #4, 0.0) + #3"
+        assert graph.pv_formula(None) == "COALESCE(#4, #2, 0.0)"