From 013eeaab6a1b6eb4bf4ee72ea271d08bf7e7b729 Mon Sep 17 00:00:00 2001
From: Javier Viola <javier@parity.io>
Date: Tue, 26 Nov 2024 12:07:04 +0100
Subject: [PATCH] fix #274, refactor trait and make 'nodes' public again

---
 .../orchestrator/src/network/chain_upgrade.rs | 33 ++++--------------
 crates/orchestrator/src/network/parachain.rs  | 30 ++++++++++++++--
 crates/orchestrator/src/network/relaychain.rs | 34 ++++++++++++++++---
 3 files changed, 64 insertions(+), 33 deletions(-)

diff --git a/crates/orchestrator/src/network/chain_upgrade.rs b/crates/orchestrator/src/network/chain_upgrade.rs
index 4da709f39..18ef9c6a8 100644
--- a/crates/orchestrator/src/network/chain_upgrade.rs
+++ b/crates/orchestrator/src/network/chain_upgrade.rs
@@ -9,37 +9,18 @@ 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"))?
diff --git a/crates/orchestrator/src/network/parachain.rs b/crates/orchestrator/src/network/parachain.rs
index 6d5cd38d9..14abbc59a 100644
--- a/crates/orchestrator/src/network/parachain.rs
+++ b/crates/orchestrator/src/network/parachain.rs
@@ -4,6 +4,7 @@ use std::{
 };
 
 use async_trait::async_trait;
+use anyhow::anyhow;
 use provider::types::TransferedFile;
 use serde::Serialize;
 use subxt::{dynamic::Value, tx::TxStatus, OnlineClient, SubstrateConfig};
@@ -13,7 +14,7 @@ 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 +30,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
     }
 }
 
diff --git a/crates/orchestrator/src/network/relaychain.rs b/crates/orchestrator/src/network/relaychain.rs
index 5c4c52bbc..1b2f99129 100644
--- a/crates/orchestrator/src/network/relaychain.rs
+++ b/crates/orchestrator/src/network/relaychain.rs
@@ -1,10 +1,10 @@
 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 +16,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 +55,9 @@ impl Relaychain {
     }
 
     // Public API
+    pub fn nodes(&self) -> Vec<&NetworkNode> {
+        self.nodes.iter().collect()
+    }
 
     /// Get chain name
     pub fn chain(&self) -> &str {