Feat upgrade runtime for parachains (#269)

- Added
   - ChainUpgrade trait, implemented for both relay and parachain.
 
- Add example of parachain upgrade.

cc @ordian

crates/examples/examples/para_upgrade.rs

@@ -0,0 +1,105 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: Apache-2.0
+
+use std::env;
+
+use anyhow::anyhow;
+use zombienet_sdk::{
+    tx_helper::{ChainUpgrade, RuntimeUpgradeOptions},
+    NetworkConfigBuilder,
+};
+
+const BEST_BLOCK_METRIC: &str = "block_height{status=\"best\"}";
+
+#[tokio::main]
+async fn main() -> Result<(), anyhow::Error> {
+    tracing_subscriber::fmt::init();
+
+    // allow to pass the upgrade path through first cli argument
+    let args: Vec<_> = env::args().collect();
+
+    let images = zombienet_sdk::environment::get_images_from_env();
+    let config = NetworkConfigBuilder::new()
+        .with_relaychain(|r| {
+            r.with_chain("rococo-local")
+                .with_default_command("polkadot")
+                .with_default_image(images.polkadot.as_str())
+                .with_node(|node| node.with_name("alice"))
+                .with_node(|node| node.with_name("bob"))
+        })
+        .with_parachain(|p| {
+            p.with_id(100)
+                .with_default_command("test-parachain")
+                .with_default_image(images.cumulus.as_str())
+                .with_collator(|c| c.with_name("collator"))
+        })
+        .build()
+        .map_err(|e| {
+            let errs = e
+                .into_iter()
+                .map(|e| e.to_string())
+                .collect::<Vec<_>>()
+                .join(" ");
+            anyhow!("config errs: {errs}")
+        })?;
+
+    let spawn_fn = zombienet_sdk::environment::get_spawn_fn();
+    let network = spawn_fn(config).await?;
+
+    // wait 2 blocks
+    let alice = network.get_node("alice")?;
+    assert!(alice
+        .wait_metric(BEST_BLOCK_METRIC, |b| b > 2_f64)
+        .await
+        .is_ok());
+
+    // get current runtime spec
+    let client = network
+        .get_node("collator")?
+        .client::<subxt::PolkadotConfig>()
+        .await?;
+    let current_runtime = client.backend().current_runtime_version().await?;
+    println!(
+        "current_runtime spec version: {:?}",
+        current_runtime.spec_version
+    );
+
+    // get current best
+    let best_block = alice.reports(BEST_BLOCK_METRIC).await?;
+
+    // upgrade runtime
+    let wasm = if args.len() > 1 {
+        args[1].clone()
+    } else if env::var("ZOMBIE_WASM_INCREMENTED_PATH").is_ok() {
+        env::var("ZOMBIE_WASM_INCREMENTED_PATH").unwrap()
+    } else {
+        panic!("You need to provide the PATH to the wasm file to use to upgrade, through first argument or 'ZOMBIE_WASM_INCREMENTED_PATH' env var");
+    };
+
+    println!("Perfoming upgrade from path {wasm}");
+
+    network
+        .parachain(100)
+        .expect("Invalid parachain Id")
+        .runtime_upgrade(RuntimeUpgradeOptions::new(wasm.as_str().into()))
+        .await?;
+
+    // wait 2 more blocks
+    alice
+        .wait_metric(BEST_BLOCK_METRIC, |x| x > best_block + 2_f64)
+        .await?;
+
+    let incremented_runtime = client.backend().current_runtime_version().await?;
+    println!(
+        "incremented_runtime spec version: {}",
+        incremented_runtime.spec_version
+    );
+
+    assert_eq!(
+        incremented_runtime.spec_version,
+        current_runtime.spec_version + 1000,
+        "version should be incremented"
+    );
+
+    Ok(())
+}

crates/orchestrator/src/network.rs

@@ -1,3 +1,4 @@
+pub mod chain_upgrade;
 pub mod node;
 pub mod parachain;
 pub mod relaychain;
@@ -563,11 +564,11 @@ impl<T: FileSystem> Network<T> {
         self.parachains.insert(para.para_id, para);
     }
 
-    pub(crate) fn name(&self) -> &str {
+    pub fn name(&self) -> &str {
         self.ns.name()
     }
 
-    pub(crate) fn parachain(&self, para_id: u32) -> Option<&Parachain> {
+    pub fn parachain(&self, para_id: u32) -> Option<&Parachain> {
         self.parachains.get(&para_id)
     }
 

crates/orchestrator/src/network/chain_upgrade.rs

@@ -0,0 +1,57 @@
+use std::str::FromStr;
+
+use anyhow::anyhow;
+use async_trait::async_trait;
+use subxt_signer::{sr25519::Keypair, SecretUri};
+
+use super::node::NetworkNode;
+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!"));
+            }
+        };
+
+        let sudo = if let Some(possible_seed) = options.seed {
+            Keypair::from_secret_key(possible_seed)
+                .map_err(|_| anyhow!("seed should return a Keypair"))?
+        } else {
+            let uri = SecretUri::from_str("//Alice")?;
+            Keypair::from_uri(&uri).map_err(|_| anyhow!("'//Alice' should return a Keypair"))?
+        };
+
+        let wasm_data = options.wasm.get_asset().await?;
+
+        tx_helper::runtime_upgrade::upgrade(node, &wasm_data, &sudo).await?;
+
+        Ok(())
+    }
+}

crates/orchestrator/src/network/parachain.rs

@@ -3,16 +3,15 @@ use std::{
     str::FromStr,
 };
 
+use async_trait::async_trait;
 use provider::types::TransferedFile;
 use serde::Serialize;
 use subxt::{dynamic::Value, tx::TxStatus, OnlineClient, SubstrateConfig};
 use subxt_signer::{sr25519::Keypair, SecretUri};
 use support::{constants::THIS_IS_A_BUG, fs::FileSystem, net::wait_ws_ready};
 use tracing::info;
 
-// use crate::generators::key::generate_pair;
-// use sp_core::{sr25519, Pair};
-use super::node::NetworkNode;
+use super::{chain_upgrade::ChainUpgrade, node::NetworkNode};
 use crate::{
     network_spec::parachain::ParachainSpec, shared::types::RegisterParachainOptions,
     ScopedFilesystem,
@@ -28,6 +27,13 @@ pub struct Parachain {
     pub(crate) files_to_inject: Vec<TransferedFile>,
 }
 
+#[async_trait]
+impl ChainUpgrade for Parachain {
+    fn nodes(&self) -> Vec<&NetworkNode> {
+        self.collators.iter().collect()
+    }
+}
+
 impl Parachain {
     pub(crate) fn new(para_id: u32) -> Self {
         Self {

crates/orchestrator/src/network/relaychain.rs

@@ -1,11 +1,10 @@
-use std::{path::PathBuf, str::FromStr};
+use std::path::PathBuf;
 
-use anyhow::anyhow;
+use async_trait::async_trait;
 use serde::Serialize;
-use subxt_signer::{sr25519::Keypair, SecretUri};
 
 use super::node::NetworkNode;
-use crate::{shared::types::RuntimeUpgradeOptions, tx_helper};
+use crate::network::chain_upgrade::ChainUpgrade;
 
 #[derive(Debug, Serialize)]
 pub struct Relaychain {
@@ -15,6 +14,13 @@ pub struct Relaychain {
     pub(crate) nodes: Vec<NetworkNode>,
 }
 
+#[async_trait]
+impl ChainUpgrade for Relaychain {
+    fn nodes(&self) -> Vec<&NetworkNode> {
+        self.nodes.iter().collect()
+    }
+}
+
 impl Relaychain {
     pub(crate) fn new(chain: String, chain_id: String, chain_spec_path: PathBuf) -> Self {
         Self {
@@ -26,54 +32,9 @@ impl Relaychain {
     }
 
     // Public API
+
+    /// Get chain name
     pub fn chain(&self) -> &str {
         &self.chain
     }
-
-    pub fn nodes(&self) -> Vec<&NetworkNode> {
-        self.nodes.iter().collect()
-    }
-
-    /// Perform a runtime upgrade (with sudo)
-    ///
-    /// This call 'System.set_code_without_checks' wrapped in
-    /// 'Sudo.sudo_unchecked_weight'
-    pub 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 relaychain", node_name));
-            }
-        } else {
-            // take the first node
-            if let Some(node) = self.nodes().first() {
-                node
-            } else {
-                return Err(anyhow!("Relaychain doesn't have any node!"));
-            }
-        };
-
-        let sudo = if let Some(possible_seed) = options.seed {
-            Keypair::from_secret_key(possible_seed)
-                .map_err(|_| anyhow!("seed should return a Keypair"))?
-        } else {
-            let uri = SecretUri::from_str("//Alice")?;
-            Keypair::from_uri(&uri).map_err(|_| anyhow!("'//Alice' should return a Keypair"))?
-        };
-
-        let wasm_data = options.wasm.get_asset().await?;
-
-        tx_helper::runtime_upgrade::upgrade(node, &wasm_data, &sudo).await?;
-
-        Ok(())
-    }
 }

crates/sdk/src/lib.rs

@@ -10,7 +10,9 @@ pub use orchestrator::{
 
 // Helpers used for interact with the network
 pub mod tx_helper {
-    pub use orchestrator::shared::types::RuntimeUpgradeOptions;
+    pub use orchestrator::{
+        network::chain_upgrade::ChainUpgrade, shared::types::RuntimeUpgradeOptions,
+    };
 }
 
 use provider::{DockerProvider, KubernetesProvider, NativeProvider};