Add tests for the logical meter (#13)

Author: shsms

Cargo.toml

@@ -13,6 +13,7 @@ chrono = "0.4"
 frequenz-microgrid-component-graph = { git = "https://github.com/frequenz-floss/frequenz-microgrid-component-graph-rs.git", rev = "0b28f99" }
 frequenz-microgrid-formula-engine = { git = "https://github.com/frequenz-floss/frequenz-microgrid-formula-engine-rs.git", rev = "463414a" }
 frequenz-resampling = { git = "https://github.com/frequenz-floss/frequenz-resampling-rs.git", rev = "ce84d66" }
+futures = "0.3.31"
 prost = "0.13"
 prost-types = "0.13"
 tokio = { version = "1.45", features = ["rt", "rt-multi-thread"] }
@@ -22,6 +23,9 @@ tracing-subscriber = { version = "0.3" }
 
 [dev-dependencies]
 tracing-subscriber = { version = "0.3", features = ["std", "env-filter"] }
+tokio = { version = "1.45", features = ["test-util"] }
+tokio-stream = { version = "0.1.17", features = ["sync"] }
+
 
 [build-dependencies]
 prost-build = { version = "0.13", features = ["cleanup-markdown"] }

examples/logical_meter.rs

@@ -18,7 +18,7 @@ async fn main() -> Result<(), Error> {
         .with(fmt::layer().with_file(true).with_line_number(true))
         .init();
 
-    let client = MicrogridClientHandle::new("http://[::1]:8800");
+    let client = MicrogridClientHandle::try_new("http://[::1]:8800").await?;
     let mut logical_meter = LogicalMeterHandle::try_new(
         client,
         LogicalMeterConfig {

src/client.rs

@@ -7,5 +7,11 @@ mod instruction;
 mod microgrid_client_actor;
 mod retry_tracker;
 
+mod microgrid_api_client;
+pub use microgrid_api_client::MicrogridApiClient;
+
 mod microgrid_client_handle;
 pub use microgrid_client_handle::MicrogridClientHandle;
+
+#[cfg(test)]
+pub(crate) mod test_utils;

src/client/microgrid_api_client.rs

@@ -0,0 +1,83 @@
+// License: MIT
+// Copyright © 2025 Frequenz Energy-as-a-Service GmbH
+
+//! The microgrid API client trait and its implementation for the generated gRPC
+//! client.
+
+use tonic::transport::Channel;
+
+use crate::proto::microgrid::v1alpha18::{
+    ListElectricalComponentConnectionsRequest, ListElectricalComponentConnectionsResponse,
+    ListElectricalComponentsRequest, ListElectricalComponentsResponse,
+    ReceiveElectricalComponentTelemetryStreamRequest,
+    ReceiveElectricalComponentTelemetryStreamResponse,
+};
+
+/// A trait representing the microgrid API client.
+///
+/// Includes all the client methods that are used by the `MicrogridClientActor`.
+#[async_trait::async_trait]
+pub trait MicrogridApiClient: Send + Sync + 'static {
+    async fn list_electrical_components(
+        &mut self,
+        request: impl tonic::IntoRequest<ListElectricalComponentsRequest> + Send,
+    ) -> std::result::Result<tonic::Response<ListElectricalComponentsResponse>, tonic::Status>;
+
+    async fn list_electrical_component_connections(
+        &mut self,
+        request: impl tonic::IntoRequest<ListElectricalComponentConnectionsRequest> + Send,
+    ) -> std::result::Result<
+        tonic::Response<ListElectricalComponentConnectionsResponse>,
+        tonic::Status,
+    >;
+
+    type TelemetryStream: futures::Stream<
+            Item = std::result::Result<
+                ReceiveElectricalComponentTelemetryStreamResponse,
+                tonic::Status,
+            >,
+        > + Send
+        + Unpin
+        + 'static;
+
+    async fn receive_electrical_component_telemetry_stream(
+        &mut self,
+        request: impl tonic::IntoRequest<ReceiveElectricalComponentTelemetryStreamRequest> + Send,
+    ) -> std::result::Result<tonic::Response<Self::TelemetryStream>, tonic::Status>;
+}
+
+/// Implement the MicrogridApiClient trait for the generated gRPC client.
+///
+/// Forwards calls to the underlying gRPC client methods, without any additional logic.
+#[async_trait::async_trait]
+impl MicrogridApiClient
+    for crate::proto::microgrid::v1alpha18::microgrid_client::MicrogridClient<Channel>
+{
+    async fn list_electrical_components(
+        &mut self,
+        request: impl tonic::IntoRequest<ListElectricalComponentsRequest> + Send,
+    ) -> std::result::Result<tonic::Response<ListElectricalComponentsResponse>, tonic::Status> {
+        self.list_electrical_components(request).await
+    }
+
+    async fn list_electrical_component_connections(
+        &mut self,
+        request: impl tonic::IntoRequest<ListElectricalComponentConnectionsRequest> + Send,
+    ) -> std::result::Result<
+        tonic::Response<ListElectricalComponentConnectionsResponse>,
+        tonic::Status,
+    > {
+        self.list_electrical_component_connections(request).await
+    }
+
+    type TelemetryStream =
+        tonic::codec::Streaming<ReceiveElectricalComponentTelemetryStreamResponse>;
+
+    async fn receive_electrical_component_telemetry_stream(
+        &mut self,
+        request: impl tonic::IntoRequest<ReceiveElectricalComponentTelemetryStreamRequest> + Send,
+    ) -> std::result::Result<tonic::Response<Self::TelemetryStream>, tonic::Status> {
+        self.receive_electrical_component_telemetry_stream(request)
+            .await
+    }
+}

src/client/microgrid_client_actor.rs

@@ -4,7 +4,7 @@
 //! The microgrid client actor that handles communication with the microgrid API.
 
 use crate::{
-    client::{instruction::Instruction, retry_tracker::RetryTracker},
+    client::{MicrogridApiClient, instruction::Instruction, retry_tracker::RetryTracker},
     proto::microgrid::v1alpha18::{
         ListElectricalComponentConnectionsRequest, ListElectricalComponentsRequest,
         ReceiveElectricalComponentTelemetryStreamRequest,
@@ -13,19 +13,15 @@ use crate::{
 };
 use std::collections::HashMap;
 
+use futures::{Stream, StreamExt};
 use tokio::{
     select,
     sync::{broadcast, mpsc},
 };
-use tonic::transport::Channel;
 use tracing::Instrument as _;
 
 use crate::{
-    Error,
-    proto::{
-        common::v1alpha8::microgrid::electrical_components::ElectricalComponentTelemetry,
-        microgrid::v1alpha18::microgrid_client::MicrogridClient,
-    },
+    Error, proto::common::v1alpha8::microgrid::electrical_components::ElectricalComponentTelemetry,
 };
 
 enum StreamStatus {
@@ -39,28 +35,20 @@ enum StreamStatus {
 ///
 /// It allows there to be multiple `MicrogridClientHandle` instances, all
 /// sharing the same connection to the microgrid API.
-pub(super) struct MicrogridClientActor {
-    url: String,
+pub(super) struct MicrogridClientActor<T> {
+    client: T,
     instructions_rx: mpsc::Receiver<Instruction>,
 }
 
-impl MicrogridClientActor {
-    pub(super) fn new(url: String, instructions_rx: mpsc::Receiver<Instruction>) -> Self {
+impl<T: MicrogridApiClient> MicrogridClientActor<T> {
+    pub(super) fn new_from_client(client: T, instructions_rx: mpsc::Receiver<Instruction>) -> Self {
         Self {
-            url,
+            client,
             instructions_rx,
         }
     }
 
     pub(super) async fn run(mut self) {
-        let mut client = match MicrogridClient::<Channel>::connect(self.url).await {
-            Ok(t) => t,
-            Err(e) => {
-                tracing::error!("Could not connect to server: {e}");
-                return;
-            }
-        };
-
         let mut component_streams: HashMap<u64, broadcast::Sender<ElectricalComponentTelemetry>> =
             HashMap::new();
 
@@ -73,7 +61,7 @@ impl MicrogridClientActor {
             select! {
                 instruction = self.instructions_rx.recv() => {
                     if let Err(e) = handle_instruction(
-                        &mut client,
+                        &mut self.client,
                         &mut component_streams,
                         instruction,
                         stream_status_tx.clone(),
@@ -102,7 +90,7 @@ impl MicrogridClientActor {
                 }
                 now = retry_timer.tick() => {
                     if let Err(e) = handle_retry_timer(
-                        &mut client,
+                        &mut self.client,
                         &mut component_streams,
                         &mut components_to_retry,
                         stream_status_tx.clone(),
@@ -117,8 +105,8 @@ impl MicrogridClientActor {
 }
 
 /// Handles the instructions received from the `MicrogridClientHandle` instances.
-async fn handle_instruction(
-    client: &mut MicrogridClient<Channel>,
+async fn handle_instruction<T: MicrogridApiClient>(
+    client: &mut T,
     component_streams: &mut HashMap<u64, broadcast::Sender<ElectricalComponentTelemetry>>,
     instruction: Option<Instruction>,
     stream_status_tx: mpsc::Sender<StreamStatus>,
@@ -199,8 +187,8 @@ async fn handle_instruction(
 
 /// Handles the retry timer, checking if the data streams for any components
 /// need to be retried and restarting their streaming tasks if necessary.
-async fn handle_retry_timer(
-    client: &mut MicrogridClient<Channel>,
+async fn handle_retry_timer<T: MicrogridApiClient>(
+    client: &mut T,
     component_streams: &mut HashMap<u64, broadcast::Sender<ElectricalComponentTelemetry>>,
     components_to_retry: &mut HashMap<u64, RetryTracker>,
     stream_status_tx: mpsc::Sender<StreamStatus>,
@@ -234,8 +222,8 @@ async fn handle_retry_timer(
 
 /// Creates a new data stream for the given component ID and starts a task to
 /// fetch data from it in a loop.
-async fn start_electrical_component_telemetry_stream(
-    client: &mut MicrogridClient<Channel>,
+async fn start_electrical_component_telemetry_stream<T: MicrogridApiClient>(
+    client: &mut T,
     electrical_component_id: u64,
     tx: broadcast::Sender<ElectricalComponentTelemetry>,
     stream_status_tx: mpsc::Sender<StreamStatus>,
@@ -283,7 +271,9 @@ async fn start_electrical_component_telemetry_stream(
 }
 
 async fn run_electrical_component_telemetry_stream(
-    mut stream: tonic::Streaming<ReceiveElectricalComponentTelemetryStreamResponse>,
+    mut stream: impl Stream<
+        Item = Result<ReceiveElectricalComponentTelemetryStreamResponse, tonic::Status>,
+    > + Unpin,
     electrical_component_id: u64,
     tx: broadcast::Sender<ElectricalComponentTelemetry>,
     stream_status_tx: mpsc::Sender<StreamStatus>,
@@ -306,30 +296,30 @@ async fn run_electrical_component_telemetry_stream(
                 });
             return;
         }
-        let message = match stream.message().await {
-            Ok(m) => m,
-            Err(e) => {
+        let message = match stream.next().await {
+            Some(m) => m,
+            None => {
                 tracing::error!(
-                    "get_component_data stream failed for {:?}: {:?}",
+                    "get_component_data stream failed for {}",
                     electrical_component_id,
-                    e
                 );
                 break;
             }
         };
         let data = match message {
-            Some(ReceiveElectricalComponentTelemetryStreamResponse { telemetry: Some(d) }) => d,
-            Some(ReceiveElectricalComponentTelemetryStreamResponse { telemetry: None }) => {
+            Ok(ReceiveElectricalComponentTelemetryStreamResponse { telemetry: Some(d) }) => d,
+            Ok(ReceiveElectricalComponentTelemetryStreamResponse { telemetry: None }) => {
                 tracing::warn!(
                     "get_component_data stream returned empty data for {}",
                     electrical_component_id
                 );
                 continue;
             }
-            None => {
+            Err(e) => {
                 tracing::warn!(
-                    "get_component_data stream ended for {:?}",
-                    electrical_component_id
+                    "get_component_data stream ended for {}: {:?}",
+                    electrical_component_id,
+                    e
                 );
                 break;
             }