Confirm components are batteries when creating a battery pool

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

Author: shsms

examples/bounds.rs

@@ -23,7 +23,7 @@ async fn main() -> Result<(), Error> {
     )
     .await?;
 
-    let mut battery_pool = microgrid.battery_pool(None);
+    let mut battery_pool = microgrid.battery_pool(None)?;
     let mut bounds_rx = battery_pool.power_bounds();
 
     while let Ok(bounds) = bounds_rx.recv().await {

src/error.rs

@@ -59,6 +59,7 @@ ErrorKind!(
     (ChronoError, chrono_error),
     (DroppedUnusedFormulas, dropped_unused_formulas),
     (FormulaEngineError, formula_engine_error),
+    (InvalidComponent, invalid_component),
     (Internal, internal),
     (APIServerError, api_server_error),
 );

src/microgrid.rs

@@ -58,8 +58,8 @@ impl Microgrid {
         self.logical_meter.clone()
     }
 
-    pub fn battery_pool(&self, component_ids: Option<Vec<u64>>) -> BatteryPool {
-        BatteryPool::new(
+    pub fn battery_pool(&self, component_ids: Option<Vec<u64>>) -> Result<BatteryPool, Error> {
+        BatteryPool::try_new(
             component_ids.map(|ids| ids.into_iter().collect()),
             self.client.clone(),
             self.logical_meter.clone(),

src/microgrid/battery_pool.rs

@@ -36,30 +36,48 @@ pub struct BatteryPool {
 impl BatteryPool {
     /// Creates a new `BatteryPool` instance with the given component IDs,
     /// client and logical meter handles.
-    pub(crate) fn new(
+    pub(crate) fn try_new(
         component_ids: Option<BTreeSet<u64>>,
         client: MicrogridClientHandle,
         logical_meter: LogicalMeterHandle,
-    ) -> Self {
-        Self {
+    ) -> Result<Self, Error> {
+        let this = Self {
             component_ids,
             client,
             logical_meter,
             snapshot_tx: None,
             bounds_tx: None,
+        };
+        if let Some(ids) = &this.component_ids {
+            if ids.is_empty() {
+                let e = "component_ids cannot be an empty set".to_string();
+                tracing::error!("{e}");
+                return Err(Error::invalid_component(e));
+            }
+            // Validate that all provided IDs correspond to batteries in the graph.
+            if !ids.is_subset(&this.get_all_battery_ids()) {
+                let e = format!("All component_ids {:?} must be batteries.", ids);
+                tracing::error!("{e}");
+                return Err(Error::invalid_component(e));
+            }
         }
+        Ok(this)
+    }
+
+    fn get_all_battery_ids(&self) -> BTreeSet<u64> {
+        self.logical_meter
+            .graph()
+            .components()
+            .filter(|c| c.category() == ElectricalComponentCategory::Battery)
+            .map(|c| c.id)
+            .collect()
     }
 
     pub(crate) fn get_battery_ids(&self) -> BTreeSet<u64> {
         if let Some(ids) = &self.component_ids {
             ids.clone()
         } else {
-            self.logical_meter
-                .graph()
-                .components()
-                .filter(|c| c.category() == ElectricalComponentCategory::Battery)
-                .map(|c| c.id)
-                .collect()
+            self.get_all_battery_ids()
         }
     }
 

src/microgrid/telemetry_tracker/battery_pool_telemetry_tracker.rs

@@ -242,7 +242,7 @@ mod tests {
         )
         .await
         .unwrap();
-        BatteryPool::new(None, client, lm)
+        BatteryPool::try_new(None, client, lm).unwrap()
     }
 
     /// Drains `rx` for up to `steps` * 100ms of simulated time, returning the