fix #274, refactor trait and make 'nodes' public again

crates/orchestrator/src/network/chain_upgrade.rs

@@ -9,37 +9,22 @@ use crate::{shared::types::RuntimeUpgradeOptions, tx_helper};
 
 #[async_trait]
 pub trait ChainUpgrade {
-    /// Get a vec of [`NetworkNode`]
-    fn nodes(&self) -> Vec<&NetworkNode>;
-
     /// Perform a runtime upgrade (with sudo)
     ///
     /// This call 'System.set_code_without_checks' wrapped in
     /// 'Sudo.sudo_unchecked_weight'
-    async fn runtime_upgrade(&self, options: RuntimeUpgradeOptions) -> Result<(), anyhow::Error> {
-        // check if the node is valid first
-        let node = if let Some(node_name) = options.node_name {
-            if let Some(node) = self
-                .nodes()
-                .into_iter()
-                .find(|node| node.name() == node_name)
-            {
-                node
-            } else {
-                return Err(anyhow!(
-                    "Node: {} is not part of the set of nodes",
-                    node_name
-                ));
-            }
-        } else {
-            // take the first node
-            if let Some(node) = self.nodes().first() {
-                node
-            } else {
-                return Err(anyhow!("chain doesn't have any node!"));
-            }
-        };
 
+    async fn runtime_upgrade(&self, options: RuntimeUpgradeOptions) -> Result<(), anyhow::Error>;
+
+    /// Perform a runtime upgrade (with sudo), inner call with the node pass as arg.
+    ///
+    /// This call 'System.set_code_without_checks' wrapped in
+    /// 'Sudo.sudo_unchecked_weight'
+    async fn perform_runtime_upgrade(
+        &self,
+        node: &NetworkNode,
+        options: RuntimeUpgradeOptions,
+    ) -> Result<(), anyhow::Error> {
         let sudo = if let Some(possible_seed) = options.seed {
             Keypair::from_secret_key(possible_seed)
                 .map_err(|_| anyhow!("seed should return a Keypair"))?

crates/orchestrator/src/network/parachain.rs

@@ -3,6 +3,7 @@ use std::{
     str::FromStr,
 };
 
+use anyhow::anyhow;
 use async_trait::async_trait;
 use provider::types::TransferedFile;
 use serde::Serialize;
@@ -13,7 +14,8 @@ use tracing::info;
 
 use super::{chain_upgrade::ChainUpgrade, node::NetworkNode};
 use crate::{
-    network_spec::parachain::ParachainSpec, shared::types::RegisterParachainOptions,
+    network_spec::parachain::ParachainSpec,
+    shared::types::{RegisterParachainOptions, RuntimeUpgradeOptions},
     ScopedFilesystem,
 };
 
@@ -29,8 +31,31 @@ pub struct Parachain {
 
 #[async_trait]
 impl ChainUpgrade for Parachain {
-    fn nodes(&self) -> Vec<&NetworkNode> {
-        self.collators.iter().collect()
+    async fn runtime_upgrade(&self, options: RuntimeUpgradeOptions) -> Result<(), anyhow::Error> {
+        // check if the node is valid first
+        let node = if let Some(node_name) = &options.node_name {
+            if let Some(node) = self
+                .collators()
+                .into_iter()
+                .find(|node| node.name() == node_name)
+            {
+                node
+            } else {
+                return Err(anyhow!(
+                    "Node: {} is not part of the set of nodes",
+                    node_name
+                ));
+            }
+        } else {
+            // take the first node
+            if let Some(node) = self.collators().first() {
+                node
+            } else {
+                return Err(anyhow!("chain doesn't have any node!"));
+            }
+        };
+
+        self.perform_runtime_upgrade(node, options).await
     }
 }
 

crates/orchestrator/src/network/relaychain.rs

@@ -1,10 +1,11 @@
 use std::path::PathBuf;
 
+use anyhow::anyhow;
 use async_trait::async_trait;
 use serde::Serialize;
 
 use super::node::NetworkNode;
-use crate::network::chain_upgrade::ChainUpgrade;
+use crate::{network::chain_upgrade::ChainUpgrade, shared::types::RuntimeUpgradeOptions};
 
 #[derive(Debug, Serialize)]
 pub struct Relaychain {
@@ -16,8 +17,31 @@ pub struct Relaychain {
 
 #[async_trait]
 impl ChainUpgrade for Relaychain {
-    fn nodes(&self) -> Vec<&NetworkNode> {
-        self.nodes.iter().collect()
+    async fn runtime_upgrade(&self, options: RuntimeUpgradeOptions) -> Result<(), anyhow::Error> {
+        // check if the node is valid first
+        let node = if let Some(node_name) = &options.node_name {
+            if let Some(node) = self
+                .nodes()
+                .into_iter()
+                .find(|node| node.name() == node_name)
+            {
+                node
+            } else {
+                return Err(anyhow!(
+                    "Node: {} is not part of the set of nodes",
+                    node_name
+                ));
+            }
+        } else {
+            // take the first node
+            if let Some(node) = self.nodes().first() {
+                node
+            } else {
+                return Err(anyhow!("chain doesn't have any node!"));
+            }
+        };
+
+        self.perform_runtime_upgrade(node, options).await
     }
 }
 
@@ -32,6 +56,9 @@ impl Relaychain {
     }
 
     // Public API
+    pub fn nodes(&self) -> Vec<&NetworkNode> {
+        self.nodes.iter().collect()
+    }
 
     /// Get chain name
     pub fn chain(&self) -> &str {