pezkuwi-sdk/vendor/pezkuwi-zombienet-sdk/crates/sdk/tests/smoke.rs

use std::{path::PathBuf, time::Instant};

use configuration::{NetworkConfig, NetworkConfigBuilder};
use futures::{stream::StreamExt, try_join};
use orchestrator::{AddCollatorOptions, AddNodeOptions};
use pezkuwi_zombienet_sdk::environment::{get_attach_fn, get_spawn_fn};

fn small_network() -> NetworkConfig {
	NetworkConfigBuilder::new()
		.with_relaychain(|r| {
			r.with_chain("rococo-local")
				.with_default_command("polkadot")
				.with_default_image("docker.io/parity/polkadot:v1.20.2")
				.with_validator(|node| node.with_name("alice"))
				.with_validator(|node| node.with_name("bob"))
		})
		.with_parachain(|p| {
			p.with_id(2000).cumulus_based(true).with_collator(|n| {
				n.with_name("collator")
					.with_command("polkadot-parachain")
					.with_image("docker.io/parity/polkadot-parachain:1.7.0")
			})
		})
		.with_parachain(|p| {
			p.with_id(3000).cumulus_based(true).with_collator(|n| {
				n.with_name("collator-new")
					.with_command("polkadot-parachain")
					.with_image("docker.io/parity/polkadot-parachain:v1.20.2")
			})
		})
		.with_global_settings(|g| {
			g.with_base_dir(PathBuf::from("/tmp/zombie-1")).with_tear_down_on_failure(false)
		})
		.build()
		.unwrap()
}

#[tokio::test(flavor = "multi_thread")]
async fn ci_k8s_basic_functionalities_should_works() {
	let _ = tracing_subscriber::fmt::try_init();

	const BEST_BLOCK_METRIC: &str = "block_height{status=\"best\"}";
	let now = Instant::now();

	let config = small_network();
	let spawn_fn = get_spawn_fn();

	let network = spawn_fn(config).await.unwrap();

	let elapsed = now.elapsed();
	println!("🚀🚀🚀🚀 network deployed in {elapsed:.2?}");

	// detach and attach to running
	network.detach().await;
	drop(network);
	let attach_fn = get_attach_fn();
	let zombie_path = PathBuf::from("/tmp/zombie-1/zombie.json");
	let mut network = attach_fn(zombie_path).await.unwrap();

	// Get a ref to the node
	let alice = network.get_node("alice").unwrap();

	let (_best_block_pass, client) = try_join!(
		alice.wait_metric(BEST_BLOCK_METRIC, |x| x > 5_f64),
		alice.wait_client::<pezkuwi_subxt::PezkuwiConfig>()
	)
	.unwrap();

	alice.wait_log_line_count("*rted #1*", true, 10).await.unwrap();

	// check best block through metrics with timeout
	assert!(alice
		.wait_metric_with_timeout(BEST_BLOCK_METRIC, |x| x > 10_f64, 45_u32)
		.await
		.is_ok());

	// ensure timeout error
	let best_block = alice.reports(BEST_BLOCK_METRIC).await.unwrap();
	let res = alice
		.wait_metric_with_timeout(BEST_BLOCK_METRIC, |x| x > (best_block * 2_f64), 10_u32)
		.await;

	assert!(res.is_err());

	// get single metric
	let role = alice.reports("node_roles").await.unwrap();
	println!("Role is {role}");
	assert_eq!(role, 4.0);

	// subxt
	// wait 3 blocks
	let mut blocks = client.blocks().subscribe_finalized().await.unwrap().take(3);
	while let Some(block) = blocks.next().await {
		println!("Block #{}", block.unwrap().header().number);
	}

	// drop the client
	drop(client);

	// check best block through metrics
	let best_block = alice.reports("block_height{status=\"best\"}").await.unwrap();

	assert!(best_block >= 2.0, "Current best {best_block}");

	// collator
	let collator = network.get_node("collator").unwrap();
	let client = collator.wait_client::<pezkuwi_subxt::PezkuwiConfig>().await.unwrap();

	// wait 3 blocks
	let mut blocks = client.blocks().subscribe_finalized().await.unwrap().take(3);
	while let Some(block) = blocks.next().await {
		println!("Block (para) #{}", block.unwrap().header().number);
	}

	// add node
	let opts = AddNodeOptions { rpc_port: Some(9444), is_validator: true, ..Default::default() };

	network.add_node("new1", opts).await.unwrap();

	// add collator
	let col_opts = AddCollatorOptions {
		command: Some("polkadot-parachain".try_into().unwrap()),
		image: Some("docker.io/parity/polkadot-parachain:1.7.0".try_into().unwrap()),
		..Default::default()
	};

	network.add_collator("new-col-1", col_opts, 2000).await.unwrap();

	// pause / resume
	let alice = network.get_node("alice").unwrap();
	alice.pause().await.unwrap();
	tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;

	let res_err = alice.wait_metric_with_timeout(BEST_BLOCK_METRIC, |x| x > 5_f64, 5_u32).await;

	assert!(res_err.is_err());

	alice.resume().await.unwrap();
	alice.wait_metric_with_timeout(BEST_BLOCK_METRIC, |x| x > 5_f64, 5_u32).await.unwrap();

	// timeout connecting ws
	let collator = network.get_node("collator").unwrap();
	collator.pause().await.unwrap();
	tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;

	let r = collator.wait_client_with_timeout::<pezkuwi_subxt::PezkuwiConfig>(1_u32).await;
	assert!(r.is_err());

	// tear down (optional if you don't detach the network)
	network.destroy().await.unwrap();
}