Merge pull request #6 from dknopik/better-host-config

config: better way to specify custom hosts
ethereum · Oct 27, 2024 · 676eee1 · 676eee1
2 parents 59e058b + 7a6f95e
commit 676eee1
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 18 deletions.
diff --git a/lib/src/config/shadow.rs b/lib/src/config/shadow.rs
@@ -129,7 +129,7 @@ pub struct HostsMut<'a> {
 }
 
 impl<'a> Iterator for HostsMut<'a> {
-    type Item = Result<UntypedHost<&'a mut Mapping>, Error>;
+    type Item = Result<&'a mut Mapping, Error>;
 
     fn next(&mut self) -> Option<Self::Item> {
         self.hosts
@@ -138,15 +138,6 @@ impl<'a> Iterator for HostsMut<'a> {
             .map(|(_, host)| {
                 host.as_mapping_mut()
                     .ok_or_else(|| Error::ExpectedOtherType("host".to_string()))
-                    .map(UntypedHost)
             })
     }
 }
-
-pub struct UntypedHost<T>(T);
-
-impl<'a> UntypedHost<&'a mut Mapping> {
-    pub fn network_node_id_mut(&mut self) -> Option<&mut Value> {
-        self.0.get_mut("network_node_id")
-    }
-}
diff --git a/lib/src/error.rs b/lib/src/error.rs
@@ -51,4 +51,8 @@ pub enum Error {
             validator client"
     )]
     MissingValidatorCount,
+    #[error(
+        "You must specify a reliability and location, and no network_node_id in your custom hosts"
+    )]
+    InvalidShadowHost,
 }
diff --git a/lib/src/lib.rs b/lib/src/lib.rs
@@ -110,15 +110,22 @@ pub fn generate<T: TryInto<FullConfig, Error = Error>>(
 
     // postprocessing given shadow config values: overwrite string network ids
     for host in shadow_config.hosts_mut()? {
-        if let Some(node_id) = host?.network_node_id_mut() {
-            if let Some((location, reliability)) = node_id.as_str().and_then(|s| s.split_once('-'))
-            {
-                let node = network_graph.assign_network_node(location, reliability)?;
-                *node_id = Value::Number(node.id().into());
-            } else {
-                return Err(Error::ExpectedOtherType("network_node_id".to_string()));
-            }
+        let mapping = host?;
+        if mapping.get("network_node_id").is_some() {
+            return Err(Error::InvalidShadowHost);
         }
+        let Value::String(location) = mapping.remove("location").ok_or(Error::InvalidShadowHost)?
+        else {
+            return Err(Error::ExpectedOtherType("location".to_string()));
+        };
+        let Value::String(reliability) = mapping
+            .remove("reliability")
+            .ok_or(Error::InvalidShadowHost)?
+        else {
+            return Err(Error::ExpectedOtherType("reliability".to_string()));
+        };
+        let node = network_graph.assign_network_node(&location, &reliability)?;
+        mapping.insert("network_node_id".into(), node.id().into());
     }
 
     info!("Generating nodes");