Use component data wrappers in mock microgrid

... instead of creating protobuf message and use it
to create component data.

Signed-off-by: ela-kotulska-frequenz <[email protected]>

Author: ela-kotulska-frequenz

tests/timeseries/mock_microgrid.py

@@ -7,18 +7,10 @@
 
 import asyncio
 import typing
+from datetime import datetime, timezone
 from typing import Callable, Set, Tuple
 
-from frequenz.api.microgrid import (
-    battery_pb2,
-    common_pb2,
-    ev_charger_pb2,
-    inverter_pb2,
-    meter_pb2,
-    microgrid_pb2,
-)
 from frequenz.channels import Broadcast
-from google.protobuf.timestamp_pb2 import Timestamp  # pylint: disable=no-name-in-module
 from pytest_mock import MockerFixture
 
 from frequenz.sdk import microgrid
@@ -32,17 +24,22 @@
 )
 from frequenz.sdk.microgrid.client import Connection
 from frequenz.sdk.microgrid.component import (
-    BatteryData,
     Component,
     ComponentCategory,
     ComponentData,
-    EVChargerData,
-    InverterData,
+    EVChargerCableState,
+    EVChargerComponentState,
     InverterType,
-    MeterData,
 )
 from tests.utils.mock_microgrid import MockMicrogridClient
 
+from ..utils.component_data_wrapper import (
+    BatteryDataWrapper,
+    EvChargerDataWrapper,
+    InverterDataWrapper,
+    MeterDataWrapper,
+)
+
 
 class MockMicrogrid:  # pylint: disable=too-many-instance-attributes
     """Setup a MockApi instance with multiple component layouts for tests."""
@@ -79,7 +76,8 @@ def __init__(self, grid_side_meter: bool, sample_rate_s: float = 0.01):
         self.evc_ids: list[int] = []
         self.meter_ids: list[int] = [4]
 
-        self.evc_states: dict[int, ev_charger_pb2.State] = {}
+        self.evc_component_states: dict[int, EVChargerComponentState] = {}
+        self.evc_cable_states: dict[int, EVChargerCableState] = {}
 
         self._streaming_coros: list[typing.Coroutine[None, None, None]] = []
         self._streaming_tasks: list[asyncio.Task[None]] = []
@@ -105,11 +103,10 @@ async def start(
         return ret
 
     async def _comp_data_send_task(
-        self, comp_id: int, make_comp_data: Callable[[int, Timestamp], ComponentData]
+        self, comp_id: int, make_comp_data: Callable[[int, datetime], ComponentData]
     ) -> None:
         for value in range(1, 2000):
-            timestamp = Timestamp()
-            timestamp.GetCurrentTime()
+            timestamp = datetime.now(tz=timezone.utc)
             val_to_send = value + int(comp_id / 10)
             # for inverters with component_id > 100, send only half the messages.
             if comp_id % 10 == self.inverter_id_suffix:
@@ -123,27 +120,11 @@ def _start_meter_streaming(self, meter_id: int) -> None:
         self._streaming_coros.append(
             self._comp_data_send_task(
                 meter_id,
-                lambda value, ts: MeterData.from_proto(
-                    microgrid_pb2.ComponentData(
-                        id=meter_id,
-                        ts=ts,
-                        meter=meter_pb2.Meter(
-                            data=meter_pb2.Data(
-                                ac=common_pb2.AC(
-                                    power_active=common_pb2.Metric(value=value),
-                                    phase_1=common_pb2.AC.ACPhase(
-                                        current=common_pb2.Metric(value=value + 100.0)
-                                    ),
-                                    phase_2=common_pb2.AC.ACPhase(
-                                        current=common_pb2.Metric(value=value + 101.0)
-                                    ),
-                                    phase_3=common_pb2.AC.ACPhase(
-                                        current=common_pb2.Metric(value=value + 102.0)
-                                    ),
-                                )
-                            )
-                        ),
-                    ),
+                lambda value, ts: MeterDataWrapper(
+                    component_id=meter_id,
+                    timestamp=ts,
+                    active_power=value,
+                    current_per_phase=(value + 100.0, value + 101.0, value + 102.0),
                 ),
             )
         )
@@ -152,16 +133,8 @@ def _start_battery_streaming(self, bat_id: int) -> None:
         self._streaming_coros.append(
             self._comp_data_send_task(
                 bat_id,
-                lambda value, ts: BatteryData.from_proto(
-                    microgrid_pb2.ComponentData(
-                        id=bat_id,
-                        ts=ts,
-                        battery=battery_pb2.Battery(
-                            data=battery_pb2.Data(
-                                soc=common_pb2.MetricAggregation(avg=value)
-                            )
-                        ),
-                    )
+                lambda value, ts: BatteryDataWrapper(
+                    component_id=bat_id, timestamp=ts, soc=value
                 ),
             )
         )
@@ -170,18 +143,8 @@ def _start_inverter_streaming(self, inv_id: int) -> None:
         self._streaming_coros.append(
             self._comp_data_send_task(
                 inv_id,
-                lambda value, ts: InverterData.from_proto(
-                    microgrid_pb2.ComponentData(
-                        id=inv_id,
-                        ts=ts,
-                        inverter=inverter_pb2.Inverter(
-                            data=inverter_pb2.Data(
-                                ac=common_pb2.AC(
-                                    power_active=common_pb2.Metric(value=value)
-                                )
-                            )
-                        ),
-                    )
+                lambda value, ts: InverterDataWrapper(
+                    component_id=inv_id, timestamp=ts, active_power=value
                 ),
             )
         )
@@ -190,28 +153,13 @@ def _start_ev_charger_streaming(self, evc_id: int) -> None:
         self._streaming_coros.append(
             self._comp_data_send_task(
                 evc_id,
-                lambda value, ts: EVChargerData.from_proto(
-                    microgrid_pb2.ComponentData(
-                        id=evc_id,
-                        ts=ts,
-                        ev_charger=ev_charger_pb2.EVCharger(
-                            data=ev_charger_pb2.Data(
-                                ac=common_pb2.AC(
-                                    power_active=common_pb2.Metric(value=value),
-                                    phase_1=common_pb2.AC.ACPhase(
-                                        current=common_pb2.Metric(value=value + 10.0)
-                                    ),
-                                    phase_2=common_pb2.AC.ACPhase(
-                                        current=common_pb2.Metric(value=value + 11.0)
-                                    ),
-                                    phase_3=common_pb2.AC.ACPhase(
-                                        current=common_pb2.Metric(value=value + 12.0)
-                                    ),
-                                )
-                            ),
-                            state=self.evc_states[evc_id],
-                        ),
-                    ),
+                lambda value, ts: EvChargerDataWrapper(
+                    component_id=evc_id,
+                    timestamp=ts,
+                    active_power=value,
+                    current_per_phase=(value + 10.0, value + 11.0, value + 12.0),
+                    component_state=self.evc_component_states[evc_id],
+                    cable_state=self.evc_cable_states[evc_id],
                 ),
             ),
         )
@@ -297,10 +245,9 @@ def add_ev_chargers(self, count: int) -> None:
             self._id_increment += 1
 
             self.evc_ids.append(evc_id)
-            self.evc_states[evc_id] = ev_charger_pb2.State(
-                cable_state=ev_charger_pb2.CABLE_STATE_UNPLUGGED,
-                component_state=ev_charger_pb2.COMPONENT_STATE_READY,
-            )
+            self.evc_component_states[evc_id] = EVChargerComponentState.READY
+            self.evc_cable_states[evc_id] = EVChargerCableState.UNPLUGGED
+
             self._components.add(
                 Component(
                     evc_id,

tests/timeseries/test_ev_charger_pool.py

@@ -9,11 +9,14 @@
 from math import isclose
 from typing import Optional
 
-from frequenz.api.microgrid import ev_charger_pb2
 from pytest_mock import MockerFixture
 
 from frequenz.sdk import microgrid
-from frequenz.sdk.microgrid.component import ComponentMetricId
+from frequenz.sdk.microgrid.component import (
+    ComponentMetricId,
+    EVChargerCableState,
+    EVChargerComponentState,
+)
 from frequenz.sdk.timeseries.ev_charger_pool import EVChargerPool
 from frequenz.sdk.timeseries.ev_charger_pool._state_tracker import (
     EVChargerPoolStates,
@@ -58,30 +61,24 @@ async def check_next_state(
         ## check that EV_PLUGGED state gets set
         await asyncio.sleep(0.02)
         evc_2_id = mockgrid.evc_ids[2]
-        mockgrid.evc_states[evc_2_id] = ev_charger_pb2.State(
-            component_state=ev_charger_pb2.COMPONENT_STATE_READY,
-            cable_state=ev_charger_pb2.CABLE_STATE_EV_PLUGGED,
-        )
+        mockgrid.evc_cable_states[evc_2_id] = EVChargerCableState.EV_PLUGGED
+        mockgrid.evc_component_states[evc_2_id] = EVChargerComponentState.READY
         expected_states[evc_2_id] = EVChargerState.EV_PLUGGED
         await check_next_state(expected_states, (evc_2_id, EVChargerState.EV_PLUGGED))
 
         ## check that EV_LOCKED state gets set
         await asyncio.sleep(0.03)
         evc_3_id = mockgrid.evc_ids[3]
-        mockgrid.evc_states[evc_3_id] = ev_charger_pb2.State(
-            component_state=ev_charger_pb2.COMPONENT_STATE_READY,
-            cable_state=ev_charger_pb2.CABLE_STATE_EV_LOCKED,
-        )
+        mockgrid.evc_cable_states[evc_3_id] = EVChargerCableState.EV_LOCKED
+        mockgrid.evc_component_states[evc_3_id] = EVChargerComponentState.READY
         expected_states[evc_3_id] = EVChargerState.EV_LOCKED
         await check_next_state(expected_states, (evc_3_id, EVChargerState.EV_LOCKED))
 
         ## check that ERROR state gets set
         await asyncio.sleep(0.1)
         evc_1_id = mockgrid.evc_ids[1]
-        mockgrid.evc_states[evc_1_id] = ev_charger_pb2.State(
-            component_state=ev_charger_pb2.COMPONENT_STATE_ERROR,
-            cable_state=ev_charger_pb2.CABLE_STATE_EV_LOCKED,
-        )
+        mockgrid.evc_cable_states[evc_1_id] = EVChargerCableState.EV_LOCKED
+        mockgrid.evc_component_states[evc_1_id] = EVChargerComponentState.ERROR
         expected_states[evc_1_id] = EVChargerState.ERROR
         await check_next_state(expected_states, (evc_1_id, EVChargerState.ERROR))