zombienet_orchestrator/network/
node.rs

1#![allow(clippy::result_large_err)]
2use std::{
3    collections::{HashMap, HashSet},
4    path::{Path, PathBuf},
5    sync::{
6        atomic::{AtomicU32, Ordering},
7        Arc,
8    },
9    time::Duration,
10};
11
12use anyhow::anyhow;
13use configuration::types::{Arg, AssetLocation};
14use prom_metrics_parser::MetricMap;
15use provider::{
16    types::{ExecutionResult, RunScriptOptions},
17    DynNode,
18};
19use serde::{Deserialize, Serialize, Serializer};
20use subxt::{backend::rpc::RpcClient, OnlineClient, PolkadotConfig};
21use support::net::{skip_err_while_waiting, wait_ws_ready};
22use thiserror::Error;
23use tokio::sync::RwLock;
24use tracing::{debug, trace, warn};
25
26use self::core::NodeCore;
27use crate::{
28    generators::{generate_node_command, generate_node_command_cumulus, GenCmdOptions},
29    network::NodeContext,
30    network_spec::node::NodeSpec,
31    shared::{constants::PROCESS_START_TIME_METRIC, types::NodeSnapshot},
32    tx_helper::client::get_client_from_url,
33};
34
35pub mod core;
36pub mod jam;
37#[cfg(test)]
38pub(crate) mod mock;
39pub mod process;
40pub mod spawned;
41
42pub use self::{
43    jam::JamNetworkNode,
44    process::{CustomProcessNode, CustomProcessState, ProcessPort},
45    spawned::{NodeKind, SpawnedNode},
46};
47
48type BoxedClosure = Box<dyn Fn(&str) -> Result<bool, anyhow::Error> + Send + Sync>;
49
50#[derive(Error, Debug)]
51pub enum NetworkNodeError {
52    #[error("metric '{0}' not found!")]
53    MetricNotFound(String),
54}
55
56// Log target for the internal monitor
57const MONITOR_TARGET: &str = "zombie_monitor";
58
59/// A substrate-based node (relaychain node or collator).
60///
61/// The provider handle and everything protocol-agnostic lives in
62/// [`NodeCore`]; this type adds the substrate surface on top: the subxt
63/// client, Prometheus metrics assertions and command regeneration from the
64/// [`NodeSpec`].
65#[derive(Clone, Serialize)]
66pub struct NetworkNode {
67    #[serde(flatten)]
68    pub(crate) core: NodeCore,
69    // TODO: do we need the full spec here?
70    pub(crate) spec: NodeSpec,
71    pub(crate) ws_uri: String,
72    pub(crate) internal_ws_uri: String,
73    pub(crate) multiaddr: String,
74    pub(crate) prometheus_uri: String,
75    // Store the option used to generate the cmd,
76    // since we can use it later for recalculate the (cmd, args) from a
77    // modified spec.
78    pub(crate) cmd_generator_opts: GenCmdOptions,
79    // Store the context used to generate this node.
80    pub(crate) context: NodeContext,
81    #[serde(skip)]
82    metrics_cache: Arc<RwLock<MetricMap>>,
83}
84#[derive(Deserialize)]
85pub(crate) struct RawNetworkNode {
86    pub(crate) name: String,
87    pub(crate) ws_uri: String,
88    pub(crate) internal_ws_uri: String,
89    pub(crate) prometheus_uri: String,
90    pub(crate) multiaddr: String,
91    pub(crate) spec: NodeSpec,
92    pub(crate) cmd_generator_opts: GenCmdOptions,
93    pub(crate) context: NodeContext,
94    #[serde(default)]
95    pub(crate) inner: serde_json::Value,
96}
97
98/// Result of waiting for a certain counter (number of log lines or events).
99///
100/// Indicates whether the log line count condition was met within the timeout period.
101///
102/// # Variants
103/// - `TargetReached(count)` – The predicate condition was satisfied within the timeout.
104///     * `count`: The number of matching instances at the time of satisfaction.
105/// - `TargetFailed(count)` – The condition was not met within the timeout.
106///     * `count`: The final number of matching instances at timeout expiration.
107#[derive(Debug, Clone, Copy, PartialEq, Eq)]
108pub enum WaitCount {
109    TargetReached(u32),
110    TargetFailed(u32),
111}
112
113impl WaitCount {
114    pub fn success(&self) -> bool {
115        match self {
116            Self::TargetReached(..) => true,
117            Self::TargetFailed(..) => false,
118        }
119    }
120}
121
122pub type LogLineCount = WaitCount;
123
124/// Configuration for controlling  count waiting behavior.
125///
126/// Allows specifying a custom predicate on the number of matchs,
127/// a timeout in seconds, and whether the system should wait the entire timeout duration.
128///
129/// # Fields
130/// - `predicate`: A function that takes the current value (metric, log lines) and
131///   returns `true` if the condition is satisfied.
132/// - `timeout_secs`: Maximum number of seconds to wait.
133/// - `wait_until_timeout_elapses`: If `true`, the system will continue waiting
134///   for the full timeout duration, even if the condition is already met early.
135///   Useful when you need to verify sustained absence or stability (e.g., "ensure no new logs appear").
136#[derive(Clone)]
137pub struct CountOptions {
138    pub predicate: Arc<dyn Fn(u32) -> bool + Send + Sync>,
139    pub timeout: Duration,
140    pub wait_until_timeout_elapses: bool,
141}
142
143impl CountOptions {
144    pub fn new(
145        predicate: impl Fn(u32) -> bool + 'static + Send + Sync,
146        timeout: Duration,
147        wait_until_timeout_elapses: bool,
148    ) -> Self {
149        Self {
150            predicate: Arc::new(predicate),
151            timeout,
152            wait_until_timeout_elapses,
153        }
154    }
155
156    pub fn no_occurences_within_timeout(timeout: Duration) -> Self {
157        Self::new(|n| n == 0, timeout, true)
158    }
159
160    pub fn at_least_once(timeout: Duration) -> Self {
161        Self::new(|count| count >= 1, timeout, false)
162    }
163
164    pub fn at_least(target: u32, timeout: Duration) -> Self {
165        Self::new(move |count| count >= target, timeout, false)
166    }
167
168    pub fn exactly_once(timeout: Duration) -> Self {
169        Self::new(|count| count == 1, timeout, false)
170    }
171}
172
173pub type LogLineCountOptions = CountOptions;
174
175impl NetworkNode {
176    /// Create a new NetworkNode
177    #[allow(clippy::too_many_arguments)]
178    pub(crate) fn new<T: Into<String>>(
179        name: T,
180        ws_uri: T,
181        internal_ws_uri: T,
182        prometheus_uri: T,
183        multiaddr: T,
184        spec: NodeSpec,
185        inner: DynNode,
186        cmd_generator_opts: GenCmdOptions,
187        context: NodeContext,
188    ) -> Self {
189        Self {
190            core: NodeCore::new(name, inner, NodeKind::Substrate),
191            ws_uri: ws_uri.into(),
192            internal_ws_uri: internal_ws_uri.into(),
193            prometheus_uri: prometheus_uri.into(),
194            spec,
195            cmd_generator_opts,
196            context,
197            multiaddr: multiaddr.into(),
198            metrics_cache: Arc::new(Default::default()),
199        }
200    }
201
202    /// The provider-generic part of this node.
203    pub fn core(&self) -> &NodeCore {
204        &self.core
205    }
206
207    // Delegations to `NodeCore` — kept as inherent methods so callers don't
208    // need `SpawnedNode` in scope.
209
210    /// Check if the node is currently running (not paused).
211    ///
212    /// This returns the internal running state.
213    pub fn is_running(&self) -> bool {
214        self.core.is_running()
215    }
216
217    /// Get the last timestamp when the node start.
218    pub fn last_start_ts(&self) -> u64 {
219        self.core.last_start_ts()
220    }
221
222    pub(crate) fn set_is_running(&self, is_running: bool) {
223        self.core.set_is_running(is_running);
224    }
225
226    /// Set the timestamp when the node was started
227    pub(crate) fn set_last_start_ts(&self, ts: u64) {
228        self.core.set_last_start_ts(ts);
229    }
230
231    pub fn name(&self) -> &str {
232        self.core.name()
233    }
234
235    /// Args used for bootstrap the node.
236    /// NOTE: this may not be in sync if you restart the node with [`NetworkNode::restart_with`].
237    pub fn args(&self) -> Vec<&str> {
238        self.core.args()
239    }
240
241    /// On-disk base directory of the node — root of `data/`, `relay-data/`,
242    /// `cfg/`, etc.
243    /// This will be the _base directory_ of the inner (provider) node.
244    pub fn base_dir(&self) -> &PathBuf {
245        self.core.base_dir()
246    }
247
248    /// Pause the node, this is implemented by pausing the
249    /// actual process (e.g polkadot) with sending `SIGSTOP` signal
250    ///
251    /// Note: If you're using this method with the native provider on the attached network, the live network has to be running
252    /// with global setting `teardown_on_failure` disabled.
253    pub async fn pause(&self) -> Result<(), anyhow::Error> {
254        self.core.pause().await
255    }
256
257    /// Resume the node, this is implemented by resuming the
258    /// actual process (e.g polkadot) with sending `SIGCONT` signal
259    ///
260    /// Note: If you're using this method with the native provider on the attached network, the live network has to be running
261    /// with global setting `teardown_on_failure` disabled.
262    pub async fn resume(&self) -> Result<(), anyhow::Error> {
263        self.core.resume().await
264    }
265
266    /// Restart the node using the same `cmd`, `args` and `env` (and same isolated dir)
267    ///
268    /// Note: If you're using this method with the native provider on the attached network, the live network has to be running
269    /// with global setting `teardown_on_failure` disabled.
270    pub async fn restart(&self, after: Option<Duration>) -> Result<(), anyhow::Error> {
271        self.core.restart(after).await
272    }
273
274    /// Run a script inside the node's container/environment
275    ///
276    /// The script will be uploaded to the node, made executable, and executed with
277    /// the provided arguments and environment variables.
278    ///
279    /// Returns `Ok(stdout)` on success, or `Err((exit_status, stderr))` on failure.
280    pub async fn run_script(
281        &self,
282        options: RunScriptOptions,
283    ) -> Result<ExecutionResult, anyhow::Error> {
284        self.core.run_script(options).await
285    }
286
287    /// Get the logs of the node
288    pub async fn logs(&self) -> Result<String, anyhow::Error> {
289        self.core.logs().await
290    }
291
292    /// Wait until a the number of matching log lines is reach
293    pub async fn wait_log_line_count(
294        &self,
295        pattern: impl Into<String>,
296        is_glob: bool,
297        count: usize,
298    ) -> Result<(), anyhow::Error> {
299        self.core.wait_log_line_count(pattern, is_glob, count).await
300    }
301
302    /// Waits until the number of matching log lines satisfies a custom condition,
303    /// optionally waiting for the entire duration of the timeout.
304    ///
305    /// See [`NodeCore::wait_log_line_count_with_timeout`].
306    pub async fn wait_log_line_count_with_timeout(
307        &self,
308        substring: impl Into<String>,
309        is_glob: bool,
310        options: LogLineCountOptions,
311    ) -> Result<LogLineCount, anyhow::Error> {
312        self.core
313            .wait_log_line_count_with_timeout(substring, is_glob, options)
314            .await
315    }
316
317    /// Tar the node's database into `out_path` (gzipped).
318    ///
319    /// NOTE: Currently __only__ implemented in native provider. Also,
320    /// the caller is responsible for pausing the node first;
321    /// snapshotting a running node risks a torn RocksDB state.
322    pub async fn snapshot_db(
323        &self,
324        out_path: impl AsRef<Path>,
325    ) -> Result<NodeSnapshot, anyhow::Error> {
326        let is_cumulus_based = matches!(
327            self.context,
328            NodeContext::Para {
329                is_cumulus_based: true,
330                ..
331            }
332        );
333
334        self.core.snapshot_db(out_path, is_cumulus_based).await
335    }
336
337    /// Check if the node is responsive by attempting to connect to its WebSocket endpoint.
338    ///
339    /// This performs an actual connection attempt with a short timeout (2 seconds).
340    /// Returns `true` if the node is reachable and responding, `false` otherwise.
341    ///
342    /// This is more robust than `is_running()` as it verifies the node is actually alive.
343    pub async fn is_responsive(&self) -> bool {
344        tokio::time::timeout(Duration::from_secs(2), wait_ws_ready(self.ws_uri()))
345            .await
346            .is_ok()
347    }
348
349    /// Only include the args provided by the user, filtering all
350    /// autogenerated by zombienet.
351    pub fn user_args(&self) -> Vec<String> {
352        self.spec
353            .args
354            .iter()
355            .fold(vec![], |acc, arg| [acc, arg.to_vec()].concat())
356    }
357
358    pub fn spec(&self) -> &NodeSpec {
359        &self.spec
360    }
361
362    pub fn ws_uri(&self) -> &str {
363        &self.ws_uri
364    }
365
366    pub fn multiaddr(&self) -> &str {
367        self.multiaddr.as_ref()
368    }
369
370    // Subxt
371
372    /// Get the rpc client for the node
373    pub async fn rpc(&self) -> Result<RpcClient, subxt::Error> {
374        get_client_from_url(&self.ws_uri).await
375    }
376
377    /// Get the [online client](subxt::client::OnlineClient) for the node
378    #[deprecated = "Use `wait_client` instead."]
379    pub async fn client<Config: subxt::Config>(
380        &self,
381    ) -> Result<OnlineClient<Config>, subxt::Error> {
382        self.try_client().await
383    }
384
385    /// Try to connect to the node.
386    ///
387    /// Most of the time you only want to use [`NetworkNode::wait_client`] that waits for
388    /// the node to appear before it connects to it. This function directly tries
389    /// to connect to the node and returns an error if the node is not yet available
390    /// at that point in time.
391    ///
392    /// Returns a [`OnlineClient`] on success.
393    pub async fn try_client<Config: subxt::Config>(
394        &self,
395    ) -> Result<OnlineClient<Config>, subxt::Error> {
396        get_client_from_url(&self.ws_uri).await
397    }
398
399    /// Wait until get the [online client](subxt::client::OnlineClient) for the node
400    pub async fn wait_client<Config: subxt::Config>(
401        &self,
402    ) -> Result<OnlineClient<Config>, anyhow::Error> {
403        debug!("wait_client ws_uri: {}", self.ws_uri());
404        wait_ws_ready(self.ws_uri())
405            .await
406            .map_err(|e| anyhow!("Error awaiting http_client to ws be ready, err: {e}"))?;
407
408        self.try_client()
409            .await
410            .map_err(|e| anyhow!("Can't create a subxt client, err: {e}"))
411    }
412
413    /// Wait until get the [online client](subxt::client::OnlineClient) for the node with a defined timeout
414    pub async fn wait_client_with_timeout<Config: subxt::Config>(
415        &self,
416        timeout_secs: impl Into<u64>,
417    ) -> Result<OnlineClient<Config>, anyhow::Error> {
418        debug!("waiting until subxt client is ready");
419        tokio::time::timeout(
420            Duration::from_secs(timeout_secs.into()),
421            self.wait_client::<Config>(),
422        )
423        .await?
424    }
425
426    /// Restart the node using the optional provided:
427    /// - Assets: binaries to download
428    ///   NOTE: if you want to use a diff version of polkadot you need both the
429    ///   main bin (polkadot) and the workers (polkadot-execute-worker/polkadot-prepare-worker)
430    /// - cmd: program to exec.
431    /// - args: Arguments to override the ones provided in config.
432    ///
433    /// The node will be restarted with the same `env` and isolated dirrectory (e.g same database).
434    ///
435    /// Note: If you're using this method with the native provider on the attached network, the live network has to be running
436    /// with global setting `teardown_on_failure` disabled.
437    pub async fn restart_with(
438        &self,
439        assets: Vec<AssetLocation>,
440        program: Option<String>,
441        args: Option<Vec<Arg>>,
442        after: Option<Duration>,
443    ) -> Result<(), anyhow::Error> {
444        let mut spec_cloned = self.spec.clone();
445
446        if let Some(args) = args {
447            spec_cloned.args = args;
448        }
449        if let Some(program) = program {
450            spec_cloned.command = program.as_str().try_into()?;
451        }
452
453        let (program, args) = match self.context {
454            NodeContext::Rc
455            | NodeContext::Para {
456                is_cumulus_based: false,
457                ..
458            } => generate_node_command(&spec_cloned, self.cmd_generator_opts.clone(), None),
459            NodeContext::Para {
460                para_id,
461                is_cumulus_based: true,
462            } => generate_node_command_cumulus(
463                &spec_cloned,
464                self.cmd_generator_opts.clone(),
465                para_id,
466            ),
467            NodeContext::Jam => {
468                return Err(anyhow!(
469                    "[{}] is a JAM node held as a substrate `NetworkNode`, {}",
470                    self.name(),
471                    support::constants::THIS_IS_A_BUG
472                ))
473            },
474        };
475
476        self.core
477            .restart_with(&assets, &program, &args, after)
478            .await
479    }
480
481    // Metrics assertions
482
483    /// Get metric value 'by name' from Prometheus (exposed by the node)
484    /// metric name can be:
485    /// with prefix (e.g: 'polkadot_')
486    /// with chain attribute (e.g: 'chain=rococo-local')
487    /// without prefix and/or without chain attribute
488    pub async fn reports(&self, metric_name: impl Into<String>) -> Result<f64, anyhow::Error> {
489        let metric_name = metric_name.into();
490        // force cache reload
491        self.fetch_metrics().await?;
492        // by default we treat not found as 0 (same in v1)
493        self.metric(&metric_name, true).await
494    }
495
496    /// Assert on a metric value 'by name' from Prometheus (exposed by the node)
497    /// metric name can be:
498    /// with prefix (e.g: 'polkadot_')
499    /// with chain attribute (e.g: 'chain=rococo-local')
500    /// without prefix and/or without chain attribute
501    ///
502    /// We first try to assert on the value using the cached metrics and
503    /// if not meet the criteria we reload the cache and check again
504    pub async fn assert(
505        &self,
506        metric_name: impl Into<String>,
507        value: impl Into<f64>,
508    ) -> Result<bool, anyhow::Error> {
509        let value: f64 = value.into();
510        self.assert_with(metric_name, |v| v == value).await
511    }
512
513    /// Assert on a metric value using a given predicate.
514    /// See [`NetworkNode::reports`] description for details on metric name.
515    pub async fn assert_with(
516        &self,
517        metric_name: impl Into<String>,
518        predicate: impl Fn(f64) -> bool,
519    ) -> Result<bool, anyhow::Error> {
520        let metric_name = metric_name.into();
521        // reload metrics
522        self.fetch_metrics().await?;
523        let val = self.metric(&metric_name, true).await?;
524        let log_msg = format!("🔎 Current value {val} passed to the predicated?");
525        if metric_name == PROCESS_START_TIME_METRIC {
526            trace!(target: MONITOR_TARGET, "{log_msg}");
527        } else {
528            trace!("{log_msg}");
529        }
530        Ok(predicate(val))
531    }
532
533    // Wait methods for metrics
534
535    /// Wait until a metric value pass the `predicate`
536    pub async fn wait_metric(
537        &self,
538        metric_name: impl Into<String>,
539        predicate: impl Fn(f64) -> bool,
540    ) -> Result<(), anyhow::Error> {
541        let metric_name = metric_name.into();
542        let log_msg = format!(
543            "[{}] waiting until metric {metric_name} pass the predicate",
544            self.name()
545        );
546        if metric_name == PROCESS_START_TIME_METRIC {
547            trace!(target: MONITOR_TARGET, "{log_msg}");
548        } else {
549            trace!("{log_msg}");
550        }
551
552        loop {
553            let res = self.assert_with(&metric_name, &predicate).await;
554            let log_msg = format!("res: {res:?}");
555            match res {
556                Ok(res) => {
557                    if res {
558                        return Ok(());
559                    }
560                },
561                Err(e) => match e.downcast::<reqwest::Error>() {
562                    Ok(io_err) => {
563                        if !skip_err_while_waiting(&io_err) {
564                            return Err(io_err.into());
565                        }
566                    },
567                    Err(other) => {
568                        match other.downcast::<NetworkNodeError>() {
569                            Ok(node_err) => {
570                                if !matches!(node_err, NetworkNodeError::MetricNotFound(_)) {
571                                    return Err(node_err.into());
572                                }
573                            },
574                            Err(other) => return Err(other),
575                        };
576                    },
577                },
578            }
579
580            if metric_name == PROCESS_START_TIME_METRIC {
581                trace!(target: MONITOR_TARGET, "{log_msg}");
582            } else {
583                trace!("{log_msg}");
584            }
585
586            // sleep to not spam prometheus
587            tokio::time::sleep(Duration::from_secs(1)).await;
588        }
589    }
590
591    /// Wait until a metric value pass the `predicate`
592    /// with a timeout (secs)
593    pub async fn wait_metric_with_timeout(
594        &self,
595        metric_name: impl Into<String>,
596        predicate: impl Fn(f64) -> bool,
597        timeout_secs: impl Into<u64>,
598    ) -> Result<(), anyhow::Error> {
599        let metric_name = metric_name.into();
600        let secs = timeout_secs.into();
601        let log_msg = format!(
602            "[{}] waiting until metric {metric_name} pass the predicate for {secs}s",
603            self.name()
604        );
605
606        if metric_name == PROCESS_START_TIME_METRIC {
607            trace!(target: MONITOR_TARGET, "{log_msg}");
608        } else {
609            debug!("{log_msg}");
610        }
611
612        let res = tokio::time::timeout(
613            Duration::from_secs(secs),
614            self.wait_metric(&metric_name, predicate),
615        )
616        .await;
617
618        if let Ok(inner_res) = res {
619            match inner_res {
620                Ok(_) => Ok(()),
621                Err(e) => Err(anyhow!("Error waiting for metric: {e}")),
622            }
623        } else {
624            // timeout
625            Err(anyhow!(
626                "Timeout ({secs}), waiting for metric {metric_name} pass the predicate"
627            ))
628        }
629    }
630
631    // Logs
632
633    /// Waits until the number of matching log lines satisfies a custom condition,
634    /// optionally waiting for the entire duration of the timeout.
635    ///
636    /// This method searches log lines for a given substring or glob pattern,
637    /// and evaluates the number of matching lines using a user-provided predicate function.
638    /// Optionally, it can wait for the full timeout duration to ensure the condition
639    /// holds consistently (e.g., for verifying absence of logs).
640    ///
641    /// # Arguments
642    /// * `substring` - The substring or pattern to match within log lines.
643    /// * `is_glob` - Whether to treat `substring` as a glob pattern (`true`) or a regex (`false`).
644    /// * `options` - Configuration for timeout, match count predicate, and full-duration waiting.
645    ///
646    /// # Returns
647    /// * `Ok(LogLineCount::TargetReached(n))` if the predicate was satisfied within the timeout,
648    /// * `Ok(LogLineCount::TargetFails(n))` if the predicate was not satisfied in time,
649    /// * `Err(e)` if an error occurred during log retrieval or matching.
650    ///
651    /// # Example
652    /// ```rust
653    /// # use std::{sync::Arc, time::Duration};
654    /// # use provider::NativeProvider;
655    /// # use support::{fs::local::LocalFileSystem};
656    /// # use zombienet_orchestrator::{Orchestrator, network::node::{NetworkNode, LogLineCountOptions}};
657    /// # use configuration::NetworkConfig;
658    /// # async fn example() -> Result<(), anyhow::Error> {
659    /// #   let provider = NativeProvider::new(LocalFileSystem {});
660    /// #   let orchestrator = Orchestrator::new(LocalFileSystem {}, provider);
661    /// #   let config = NetworkConfig::load_from_toml("config.toml")?;
662    /// #   let network = orchestrator.spawn(config).await?;
663    /// let node = network.get_node("alice")?;
664    /// // Wait (up to 10 seconds) until pattern occurs once
665    /// let options = LogLineCountOptions {
666    ///     predicate: Arc::new(|count| count == 1),
667    ///     timeout: Duration::from_secs(10),
668    ///     wait_until_timeout_elapses: false,
669    /// };
670    /// let result = node
671    ///     .wait_log_line_count_with_timeout("error", false, options)
672    ///     .await?;
673    /// #   Ok(())
674    /// # }
675    /// ```
676    pub async fn wait_event_count_with_timeout(
677        &self,
678        pallet: impl Into<String>,
679        variant: impl Into<String>,
680        options: CountOptions,
681    ) -> Result<WaitCount, anyhow::Error> {
682        let pallet = pallet.into();
683        let variant = variant.into();
684        debug!(
685            "waiting until match event ({pallet} {variant}) count within {} seconds",
686            options.timeout.as_secs_f64()
687        );
688
689        let init_value = Arc::new(AtomicU32::new(0));
690
691        let res = tokio::time::timeout(
692            options.timeout,
693            self.wait_event_count(&pallet, &variant, &options, init_value.clone()),
694        )
695        .await;
696
697        let q = init_value.load(Ordering::Relaxed);
698        if let Ok(inner_res) = res {
699            match inner_res {
700                Ok(_) => Ok(WaitCount::TargetReached(q)),
701                Err(e) => Err(anyhow!("Error waiting for counter: {e}")),
702            }
703        } else {
704            // timeout
705            if options.wait_until_timeout_elapses {
706                let q = init_value.load(Ordering::Relaxed);
707                if (options.predicate)(q) {
708                    Ok(LogLineCount::TargetReached(q))
709                } else {
710                    Ok(LogLineCount::TargetFailed(q))
711                }
712            } else {
713                Err(anyhow!(
714                    "Timeout ({}), waiting for counter",
715                    options.timeout.as_secs()
716                ))
717            }
718        }
719    }
720
721    //
722    async fn wait_event_count(
723        &self,
724        pallet: &str,
725        variant: &str,
726        options: &CountOptions,
727        init_count: Arc<AtomicU32>,
728    ) -> Result<(), anyhow::Error> {
729        let client: OnlineClient<PolkadotConfig> = self.wait_client().await?;
730        let mut blocks_sub: subxt::backend::StreamOf<
731            Result<
732                subxt::blocks::Block<PolkadotConfig, OnlineClient<PolkadotConfig>>,
733                subxt::Error,
734            >,
735        > = client.blocks().subscribe_finalized().await?;
736        while let Some(block) = blocks_sub.next().await {
737            let events = block?.events().await?;
738            for event in events.iter() {
739                let evt = event?;
740                if evt.pallet_name() == pallet && evt.variant_name() == variant {
741                    let old_value = init_count.fetch_add(1, Ordering::Relaxed);
742                    if !options.wait_until_timeout_elapses && (options.predicate)(old_value + 1) {
743                        return Ok(());
744                    }
745                }
746            }
747        }
748
749        Ok(())
750    }
751
752    async fn fetch_metrics(&self) -> Result<(), anyhow::Error> {
753        let response = reqwest::get(&self.prometheus_uri).await?;
754        let metrics = prom_metrics_parser::parse(&response.text().await?)?;
755        let mut cache = self.metrics_cache.write().await;
756        *cache = metrics;
757        Ok(())
758    }
759
760    /// Query individual metric by name
761    async fn metric(
762        &self,
763        metric_name: &str,
764        treat_not_found_as_zero: bool,
765    ) -> Result<f64, anyhow::Error> {
766        let mut metrics_map = self.metrics_cache.read().await;
767        if metrics_map.is_empty() {
768            // reload metrics
769            drop(metrics_map);
770            self.fetch_metrics().await?;
771            metrics_map = self.metrics_cache.read().await;
772        }
773
774        if let Some(val) = metrics_map.get(metric_name) {
775            Ok(*val)
776        } else if treat_not_found_as_zero {
777            Ok(0_f64)
778        } else {
779            Err(NetworkNodeError::MetricNotFound(metric_name.into()).into())
780        }
781    }
782
783    /// Fetches histogram buckets for a given metric from the Prometheus endpoint.
784    ///
785    /// This function retrieves histogram bucket data by parsing the Prometheus metrics
786    /// and calculating the count of observations in each bucket. It automatically appends
787    /// `_bucket` suffix to the metric name if not already present.
788    ///
789    /// # Arguments
790    /// * `metric_name` - The name of the histogram metric (with or without `_bucket` suffix)
791    /// * `label_filters` - Optional HashMap of label key-value pairs to filter metrics by
792    ///
793    /// # Returns
794    /// A HashMap where keys are the `le` bucket boundaries as strings,
795    /// and values are the count of observations in each bucket (calculated as delta from previous bucket).
796    ///
797    /// # Example
798    /// ```ignore
799    /// let buckets = node.get_histogram_buckets("polkadot_pvf_execution_time", None).await?;
800    /// // Returns: {"0.1": 5, "0.5": 10, "1.0": 3, "+Inf": 0}
801    /// ```
802    pub async fn get_histogram_buckets(
803        &self,
804        metric_name: impl AsRef<str>,
805        label_filters: Option<HashMap<String, String>>,
806    ) -> Result<HashMap<String, u64>, anyhow::Error> {
807        let metric_name = metric_name.as_ref();
808
809        // Fetch and parse metrics using the existing parser
810        let response = reqwest::get(&self.prometheus_uri).await?;
811        let metrics = prom_metrics_parser::parse(&response.text().await?)?;
812
813        // Ensure metric name has _bucket suffix
814        let resolved_metric_name = if metric_name.contains("_bucket") {
815            metric_name.to_string()
816        } else {
817            format!("{metric_name}_bucket")
818        };
819
820        // First pass: collect all matching metrics with their label counts
821        // to identify which ones have the most complete label sets
822        // Each entry contains: (full_metric_key, parsed_labels_map, cumulative_count)
823        let mut metric_entries: Vec<(String, HashMap<String, String>, u64)> = Vec::new();
824
825        for (key, &value) in metrics.iter() {
826            if !key.starts_with(&resolved_metric_name) {
827                continue;
828            }
829
830            let remaining = &key[resolved_metric_name.len()..];
831
832            let labels_str = &remaining[1..remaining.len() - 1];
833            let parsed_labels = Self::parse_label_string(labels_str);
834
835            // Must have "le" label
836            if !parsed_labels.contains_key("le") {
837                continue;
838            }
839
840            // Check if label filters match
841            if let Some(ref filters) = label_filters {
842                let mut all_match = true;
843                for (filter_key, filter_value) in filters {
844                    if parsed_labels.get(filter_key) != Some(filter_value) {
845                        all_match = false;
846                        break;
847                    }
848                }
849                if !all_match {
850                    continue;
851                }
852            }
853
854            metric_entries.push((key.clone(), parsed_labels, value as u64));
855        }
856
857        // Find the maximum number of labels (excluding "le") across all entries
858        // This helps us identify the "fullest" version of each metric
859        let max_label_count = metric_entries
860            .iter()
861            .map(|(_, labels, _)| labels.iter().filter(|(k, _)| k.as_str() != "le").count())
862            .max()
863            .unwrap_or(0);
864
865        // Second pass: collect buckets, deduplicating and preferring entries with more labels
866        let mut raw_buckets: Vec<(String, u64)> = Vec::new();
867        let mut seen_le_values = HashSet::new();
868        let mut active_series: Option<Vec<(String, String)>> = None;
869
870        for (_, parsed_labels, value) in metric_entries {
871            let le_value = parsed_labels.get("le").unwrap().clone();
872
873            // Get non-"le" labels
874            let mut non_le_labels: Vec<(String, String)> = parsed_labels
875                .iter()
876                .filter(|(k, _)| k.as_str() != "le")
877                .map(|(k, v)| (k.clone(), v.clone()))
878                .collect();
879            non_le_labels.sort();
880
881            // Only process entries that have the maximum number of labels
882            // (this filters out the parser's duplicate keys with fewer labels)
883            if non_le_labels.len() < max_label_count {
884                continue;
885            }
886
887            // Detect series changes
888            if let Some(ref prev_series) = active_series {
889                if prev_series != &non_le_labels {
890                    if !raw_buckets.is_empty() {
891                        break; // Stop at first series change
892                    }
893                    active_series = Some(non_le_labels.clone());
894                    seen_le_values.clear();
895                }
896            } else {
897                active_series = Some(non_le_labels.clone());
898            }
899
900            // Deduplicate by le value within this series
901            if !seen_le_values.insert(le_value.clone()) {
902                continue;
903            }
904
905            trace!("{} le:{} {}", resolved_metric_name, &le_value, value);
906            raw_buckets.push((le_value, value));
907        }
908
909        // Sort buckets by their "le" values
910        raw_buckets.sort_by(|a, b| Self::compare_le_values(&a.0, &b.0));
911
912        // Calculate deltas between cumulative buckets
913        let mut buckets = HashMap::new();
914        let mut previous_value = 0_u64;
915        for (le, cumulative_count) in raw_buckets {
916            if cumulative_count < previous_value {
917                warn!(
918                    "Warning: bucket count decreased from {} to {} at le={}",
919                    previous_value, cumulative_count, le
920                );
921            }
922            let delta = cumulative_count.saturating_sub(previous_value);
923            buckets.insert(le, delta);
924            previous_value = cumulative_count;
925        }
926
927        Ok(buckets)
928    }
929
930    /// Parse label string from parsed metric key.
931    ///
932    /// Takes a label string in the format `key1="value1",key2="value2"`
933    /// and returns a HashMap of key-value pairs.
934    /// Handles commas inside quoted values correctly.
935    fn parse_label_string(labels_str: &str) -> HashMap<String, String> {
936        let mut labels = HashMap::new();
937        let mut current_key = String::new();
938        let mut current_value = String::new();
939        let mut in_value = false;
940        let mut in_quotes = false;
941
942        for ch in labels_str.chars() {
943            match ch {
944                '=' if !in_quotes && !in_value => {
945                    in_value = true;
946                },
947                '"' if in_value => {
948                    in_quotes = !in_quotes;
949                },
950                ',' if !in_quotes => {
951                    // End of key-value pair
952                    if !current_key.is_empty() {
953                        labels.insert(
954                            current_key.trim().to_string(),
955                            current_value.trim().to_string(),
956                        );
957                        current_key.clear();
958                        current_value.clear();
959                        in_value = false;
960                    }
961                },
962                _ => {
963                    if in_value {
964                        current_value.push(ch);
965                    } else {
966                        current_key.push(ch);
967                    }
968                },
969            }
970        }
971
972        // Insert last pair
973        if !current_key.is_empty() {
974            labels.insert(
975                current_key.trim().to_string(),
976                current_value.trim().to_string(),
977            );
978        }
979
980        labels
981    }
982
983    /// Compare two histogram bucket boundary values for sorting.
984    ///
985    /// Treats "+Inf" as the maximum value, otherwise compares numerically.
986    fn compare_le_values(a: &str, b: &str) -> std::cmp::Ordering {
987        use std::cmp::Ordering;
988
989        // Handle +Inf specially
990        match (a, b) {
991            ("+Inf", "+Inf") => Ordering::Equal,
992            ("+Inf", _) => Ordering::Greater,
993            (_, "+Inf") => Ordering::Less,
994            _ => {
995                // Try to parse as f64 for numeric comparison
996                match (a.parse::<f64>(), b.parse::<f64>()) {
997                    (Ok(a_val), Ok(b_val)) => a_val.partial_cmp(&b_val).unwrap_or(Ordering::Equal),
998                    // Fallback to string comparison if parsing fails
999                    _ => a.cmp(b),
1000                }
1001            },
1002        }
1003    }
1004
1005    /// Waits given number of seconds until node reports that it is up and running, which
1006    /// is determined by metric 'process_start_time_seconds', which should appear,
1007    /// when node finished booting up.
1008    ///
1009    ///
1010    /// # Arguments
1011    /// * `timeout_secs` - The number of seconds to wait.
1012    ///
1013    /// # Returns
1014    /// * `Ok()` if the node is up before timeout occured.
1015    /// * `Err(e)` if timeout or other error occurred while waiting.
1016    pub async fn wait_until_is_up(
1017        &self,
1018        timeout_secs: impl Into<u64>,
1019    ) -> Result<(), anyhow::Error> {
1020        self.wait_metric_with_timeout(PROCESS_START_TIME_METRIC, |b| b >= 1.0, timeout_secs)
1021            .await
1022            .map_err(|err| anyhow::anyhow!("{}: {:?}", self.name(), err))
1023    }
1024}
1025
1026#[async_trait::async_trait]
1027impl SpawnedNode for NetworkNode {
1028    fn core(&self) -> &NodeCore {
1029        &self.core
1030    }
1031
1032    async fn wait_until_is_up(&self, timeout_secs: u64) -> Result<(), anyhow::Error> {
1033        NetworkNode::wait_until_is_up(self, timeout_secs).await
1034    }
1035}
1036
1037impl std::fmt::Debug for NetworkNode {
1038    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1039        f.debug_struct("NetworkNode")
1040            .field("inner", &"inner_skipped")
1041            .field("spec", &self.spec)
1042            .field("name", &self.name())
1043            .field("ws_uri", &self.ws_uri)
1044            .field("prometheus_uri", &self.prometheus_uri)
1045            .finish()
1046    }
1047}
1048
1049fn serialize_provider_node<S>(node: &DynNode, serializer: S) -> Result<S::Ok, S::Error>
1050where
1051    S: Serializer,
1052{
1053    erased_serde::serialize(node.as_ref(), serializer)
1054}
1055
1056// TODO: mock and impl more unit tests
1057#[cfg(test)]
1058mod tests {
1059    use std::sync::Arc;
1060
1061    use super::{mock::MockNode, *};
1062
1063    #[tokio::test(flavor = "multi_thread")]
1064    async fn test_wait_log_count_target_reached_immediately() -> Result<(), anyhow::Error> {
1065        let mock_provider = Arc::new(MockNode::new());
1066        let mock_node = NetworkNode::new(
1067            "node1",
1068            "ws_uri",
1069            "internal_ws_uri",
1070            "prometheus_uri",
1071            "multiaddr",
1072            NodeSpec::default(),
1073            mock_provider.clone(),
1074            GenCmdOptions::default(),
1075            NodeContext::Rc,
1076        );
1077
1078        mock_provider.logs_push(vec![
1079            "system booting",
1080            "stub line 1",
1081            "stub line 2",
1082            "system ready",
1083        ]);
1084
1085        // Wait (up to 10 seconds) until pattern occurs once
1086        let options = LogLineCountOptions {
1087            predicate: Arc::new(|n| n == 1),
1088            timeout: Duration::from_secs(10),
1089            wait_until_timeout_elapses: false,
1090        };
1091
1092        let log_line_count = mock_node
1093            .wait_log_line_count_with_timeout("system ready", false, options)
1094            .await?;
1095
1096        assert!(matches!(log_line_count, LogLineCount::TargetReached(1)));
1097
1098        Ok(())
1099    }
1100
1101    #[tokio::test(flavor = "multi_thread")]
1102    async fn test_wait_log_count_target_reached_after_delay() -> Result<(), anyhow::Error> {
1103        let mock_provider = Arc::new(MockNode::new());
1104        let mock_node = NetworkNode::new(
1105            "node1",
1106            "ws_uri",
1107            "internal_ws_uri",
1108            "prometheus_uri",
1109            "multiaddr",
1110            NodeSpec::default(),
1111            mock_provider.clone(),
1112            GenCmdOptions::default(),
1113            NodeContext::Rc,
1114        );
1115
1116        mock_provider.logs_push(vec![
1117            "system booting",
1118            "stub line 1",
1119            "stub line 2",
1120            "system ready",
1121        ]);
1122
1123        // Wait (up to 4 seconds) until pattern occurs twice
1124        let options = LogLineCountOptions {
1125            predicate: Arc::new(|n| n == 2),
1126            timeout: Duration::from_secs(4),
1127            wait_until_timeout_elapses: false,
1128        };
1129
1130        let task = tokio::spawn({
1131            async move {
1132                mock_node
1133                    .wait_log_line_count_with_timeout("system ready", false, options)
1134                    .await
1135                    .unwrap()
1136            }
1137        });
1138
1139        tokio::time::sleep(Duration::from_secs(2)).await;
1140
1141        mock_provider.logs_push(vec!["system ready"]);
1142
1143        let log_line_count = task.await?;
1144
1145        assert!(matches!(log_line_count, LogLineCount::TargetReached(2)));
1146
1147        Ok(())
1148    }
1149
1150    #[tokio::test(flavor = "multi_thread")]
1151    async fn test_wait_log_count_target_failed_timeout() -> Result<(), anyhow::Error> {
1152        let mock_provider = Arc::new(MockNode::new());
1153        let mock_node = NetworkNode::new(
1154            "node1",
1155            "ws_uri",
1156            "internal_ws_uri",
1157            "prometheus_uri",
1158            "multiaddr",
1159            NodeSpec::default(),
1160            mock_provider.clone(),
1161            GenCmdOptions::default(),
1162            NodeContext::Rc,
1163        );
1164
1165        mock_provider.logs_push(vec![
1166            "system booting",
1167            "stub line 1",
1168            "stub line 2",
1169            "system ready",
1170        ]);
1171
1172        // Wait (up to 2 seconds) until pattern occurs twice
1173        let options = LogLineCountOptions {
1174            predicate: Arc::new(|n| n == 2),
1175            timeout: Duration::from_secs(2),
1176            wait_until_timeout_elapses: false,
1177        };
1178
1179        let log_line_count = mock_node
1180            .wait_log_line_count_with_timeout("system ready", false, options)
1181            .await?;
1182
1183        assert!(matches!(log_line_count, LogLineCount::TargetFailed(1)));
1184
1185        Ok(())
1186    }
1187
1188    #[tokio::test(flavor = "multi_thread")]
1189    async fn test_wait_log_count_target_failed_exceeded() -> Result<(), anyhow::Error> {
1190        let mock_provider = Arc::new(MockNode::new());
1191        let mock_node = NetworkNode::new(
1192            "node1",
1193            "ws_uri",
1194            "internal_ws_uri",
1195            "prometheus_uri",
1196            "multiaddr",
1197            NodeSpec::default(),
1198            mock_provider.clone(),
1199            GenCmdOptions::default(),
1200            NodeContext::Rc,
1201        );
1202
1203        mock_provider.logs_push(vec![
1204            "system booting",
1205            "stub line 1",
1206            "stub line 2",
1207            "system ready",
1208        ]);
1209
1210        // Wait until timeout and check if pattern occurs exactly twice
1211        let options = LogLineCountOptions {
1212            predicate: Arc::new(|n| n == 2),
1213            timeout: Duration::from_secs(2),
1214            wait_until_timeout_elapses: true,
1215        };
1216
1217        let task = tokio::spawn({
1218            async move {
1219                mock_node
1220                    .wait_log_line_count_with_timeout("system ready", false, options)
1221                    .await
1222                    .unwrap()
1223            }
1224        });
1225
1226        tokio::time::sleep(Duration::from_secs(1)).await;
1227
1228        mock_provider.logs_push(vec!["system ready"]);
1229        mock_provider.logs_push(vec!["system ready"]);
1230
1231        let log_line_count = task.await?;
1232
1233        assert!(matches!(log_line_count, LogLineCount::TargetFailed(3)));
1234
1235        Ok(())
1236    }
1237
1238    #[tokio::test(flavor = "multi_thread")]
1239    async fn test_wait_log_count_target_reached_no_occurences() -> Result<(), anyhow::Error> {
1240        let mock_provider = Arc::new(MockNode::new());
1241        let mock_node = NetworkNode::new(
1242            "node1",
1243            "ws_uri",
1244            "internal_ws_uri",
1245            "prometheus_uri",
1246            "multiaddr",
1247            NodeSpec::default(),
1248            mock_provider.clone(),
1249            GenCmdOptions::default(),
1250            NodeContext::Rc,
1251        );
1252
1253        mock_provider.logs_push(vec!["system booting", "stub line 1", "stub line 2"]);
1254
1255        let task = tokio::spawn({
1256            async move {
1257                mock_node
1258                    .wait_log_line_count_with_timeout(
1259                        "system ready",
1260                        false,
1261                        // Wait until timeout and make sure pattern occurred zero times
1262                        LogLineCountOptions::no_occurences_within_timeout(Duration::from_secs(2)),
1263                    )
1264                    .await
1265                    .unwrap()
1266            }
1267        });
1268
1269        tokio::time::sleep(Duration::from_secs(1)).await;
1270
1271        mock_provider.logs_push(vec!["stub line 3"]);
1272
1273        assert!(task.await?.success());
1274
1275        Ok(())
1276    }
1277
1278    #[tokio::test(flavor = "multi_thread")]
1279    async fn test_wait_log_count_target_reached_in_range() -> Result<(), anyhow::Error> {
1280        let mock_provider = Arc::new(MockNode::new());
1281        let mock_node = NetworkNode::new(
1282            "node1",
1283            "ws_uri",
1284            "internal_ws_uri",
1285            "prometheus_uri",
1286            "multiaddr",
1287            NodeSpec::default(),
1288            mock_provider.clone(),
1289            GenCmdOptions::default(),
1290            NodeContext::Rc,
1291        );
1292
1293        mock_provider.logs_push(vec!["system booting", "stub line 1", "stub line 2"]);
1294
1295        // Wait until timeout and make sure pattern occurrence count is in range between 2 and 5
1296        let options = LogLineCountOptions {
1297            predicate: Arc::new(|n| (2..=5).contains(&n)),
1298            timeout: Duration::from_secs(2),
1299            wait_until_timeout_elapses: true,
1300        };
1301
1302        let task = tokio::spawn({
1303            async move {
1304                mock_node
1305                    .wait_log_line_count_with_timeout("system ready", false, options)
1306                    .await
1307                    .unwrap()
1308            }
1309        });
1310
1311        tokio::time::sleep(Duration::from_secs(1)).await;
1312
1313        mock_provider.logs_push(vec!["system ready", "system ready", "system ready"]);
1314
1315        assert!(task.await?.success());
1316
1317        Ok(())
1318    }
1319
1320    #[tokio::test(flavor = "multi_thread")]
1321    async fn test_wait_log_count_with_timeout_with_lookahead_regex() -> Result<(), anyhow::Error> {
1322        let mock_provider = Arc::new(MockNode::new());
1323        let mock_node = NetworkNode::new(
1324            "node1",
1325            "ws_uri",
1326            "internal_ws_uri",
1327            "prometheus_uri",
1328            "multiaddr",
1329            NodeSpec::default(),
1330            mock_provider.clone(),
1331            GenCmdOptions::default(),
1332            NodeContext::Rc,
1333        );
1334
1335        mock_provider.logs_push(vec![
1336            "system booting",
1337            "stub line 1",
1338            // this line should not match
1339            "Error importing block 0xfd66e545c446b1c01205503130b816af0ec2c0e504a8472808e6ff4a644ce1fa: block has an unknown parent",
1340            "stub line 2"
1341        ]);
1342
1343        let options = LogLineCountOptions {
1344            predicate: Arc::new(|n| n == 1),
1345            timeout: Duration::from_secs(3),
1346            wait_until_timeout_elapses: true,
1347        };
1348
1349        let task = tokio::spawn({
1350            async move {
1351                mock_node
1352                    .wait_log_line_count_with_timeout(
1353                        "error(?! importing block .*: block has an unknown parent)",
1354                        false,
1355                        options,
1356                    )
1357                    .await
1358                    .unwrap()
1359            }
1360        });
1361
1362        tokio::time::sleep(Duration::from_secs(1)).await;
1363
1364        mock_provider.logs_push(vec![
1365            "system ready",
1366            // this line should match
1367            "system error",
1368            "system ready",
1369        ]);
1370
1371        assert!(task.await?.success());
1372
1373        Ok(())
1374    }
1375
1376    #[tokio::test(flavor = "multi_thread")]
1377    async fn test_wait_log_count_with_timeout_with_lookahead_regex_fails(
1378    ) -> Result<(), anyhow::Error> {
1379        let mock_provider = Arc::new(MockNode::new());
1380        let mock_node = NetworkNode::new(
1381            "node1",
1382            "ws_uri",
1383            "internal_ws_uri",
1384            "prometheus_uri",
1385            "multiaddr",
1386            NodeSpec::default(),
1387            mock_provider.clone(),
1388            GenCmdOptions::default(),
1389            NodeContext::Rc,
1390        );
1391
1392        mock_provider.logs_push(vec![
1393            "system booting",
1394            "stub line 1",
1395            // this line should not match
1396            "Error importing block 0xfd66e545c446b1c01205503130b816af0ec2c0e504a8472808e6ff4a644ce1fa: block has an unknown parent",
1397            "stub line 2"
1398        ]);
1399
1400        let options = LogLineCountOptions {
1401            predicate: Arc::new(|n| n == 1),
1402            timeout: Duration::from_secs(6),
1403            wait_until_timeout_elapses: true,
1404        };
1405
1406        let task = tokio::spawn({
1407            async move {
1408                mock_node
1409                    .wait_log_line_count_with_timeout(
1410                        "error(?! importing block .*: block has an unknown parent)",
1411                        false,
1412                        options,
1413                    )
1414                    .await
1415                    .unwrap()
1416            }
1417        });
1418
1419        tokio::time::sleep(Duration::from_secs(1)).await;
1420
1421        mock_provider.logs_push(vec!["system ready", "system ready"]);
1422
1423        assert!(!task.await?.success());
1424
1425        Ok(())
1426    }
1427
1428    #[tokio::test(flavor = "multi_thread")]
1429    async fn test_wait_log_count_with_lockahead_regex() -> Result<(), anyhow::Error> {
1430        let mock_provider = Arc::new(MockNode::new());
1431        let mock_node = NetworkNode::new(
1432            "node1",
1433            "ws_uri",
1434            "internal_ws_uri",
1435            "prometheus_uri",
1436            "multiaddr",
1437            NodeSpec::default(),
1438            mock_provider.clone(),
1439            GenCmdOptions::default(),
1440            NodeContext::Rc,
1441        );
1442
1443        mock_provider.logs_push(vec![
1444            "system booting",
1445            "stub line 1",
1446            // this line should not match
1447            "Error importing block 0xfd66e545c446b1c01205503130b816af0ec2c0e504a8472808e6ff4a644ce1fa: block has an unknown parent",
1448            "stub line 2"
1449        ]);
1450
1451        let task = tokio::spawn({
1452            async move {
1453                mock_node
1454                    .wait_log_line_count(
1455                        "error(?! importing block .*: block has an unknown parent)",
1456                        false,
1457                        1,
1458                    )
1459                    .await
1460                    .unwrap()
1461            }
1462        });
1463
1464        tokio::time::sleep(Duration::from_secs(1)).await;
1465
1466        mock_provider.logs_push(vec![
1467            "system ready",
1468            // this line should match
1469            "system error",
1470            "system ready",
1471        ]);
1472
1473        assert!(task.await.is_ok());
1474
1475        Ok(())
1476    }
1477
1478    #[tokio::test(flavor = "multi_thread")]
1479    async fn test_wait_log_count_with_lookahead_regex_fails() -> Result<(), anyhow::Error> {
1480        let mock_provider = Arc::new(MockNode::new());
1481        let mock_node = NetworkNode::new(
1482            "node1",
1483            "ws_uri",
1484            "internal_ws_uri",
1485            "prometheus_uri",
1486            "multiaddr",
1487            NodeSpec::default(),
1488            mock_provider.clone(),
1489            GenCmdOptions::default(),
1490            NodeContext::Rc,
1491        );
1492
1493        mock_provider.logs_push(vec![
1494            "system booting",
1495            "stub line 1",
1496            // this line should not match
1497            "Error importing block 0xfd66e545c446b1c01205503130b816af0ec2c0e504a8472808e6ff4a644ce1fa: block has an unknown parent",
1498            "stub line 2"
1499        ]);
1500
1501        let options = LogLineCountOptions {
1502            predicate: Arc::new(|count| count == 1),
1503            timeout: Duration::from_secs(2),
1504            wait_until_timeout_elapses: true,
1505        };
1506
1507        let task = tokio::spawn({
1508            async move {
1509                // we expect no match, thus wait with timeout
1510                mock_node
1511                    .wait_log_line_count_with_timeout(
1512                        "error(?! importing block .*: block has an unknown parent)",
1513                        false,
1514                        options,
1515                    )
1516                    .await
1517                    .unwrap()
1518            }
1519        });
1520
1521        tokio::time::sleep(Duration::from_secs(1)).await;
1522
1523        mock_provider.logs_push(vec!["system ready", "system ready"]);
1524
1525        assert!(!task.await?.success());
1526
1527        Ok(())
1528    }
1529
1530    #[tokio::test]
1531    async fn test_get_histogram_buckets_parsing() -> Result<(), anyhow::Error> {
1532        // This test uses a mock HTTP server to simulate Prometheus metrics
1533        use std::sync::Arc;
1534
1535        // Create a mock metrics response with proper HELP and TYPE comments
1536        let mock_metrics = concat!(
1537            "# HELP substrate_block_verification_time Time taken to verify blocks\n",
1538            "# TYPE substrate_block_verification_time histogram\n",
1539            "substrate_block_verification_time_bucket{chain=\"rococo_local_testnet\",le=\"0.1\"} 10\n",
1540            "substrate_block_verification_time_bucket{chain=\"rococo_local_testnet\",le=\"0.5\"} 25\n",
1541            "substrate_block_verification_time_bucket{chain=\"rococo_local_testnet\",le=\"1.0\"} 35\n",
1542            "substrate_block_verification_time_bucket{chain=\"rococo_local_testnet\",le=\"2.5\"} 40\n",
1543            "substrate_block_verification_time_bucket{chain=\"rococo_local_testnet\",le=\"+Inf\"} 42\n",
1544            "substrate_block_verification_time_sum{chain=\"rococo_local_testnet\"} 45.5\n",
1545            "substrate_block_verification_time_count{chain=\"rococo_local_testnet\"} 42\n",
1546        );
1547
1548        // Start a mock HTTP server
1549        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?;
1550        let addr = listener.local_addr()?;
1551        let metrics = Arc::new(mock_metrics.to_string());
1552
1553        tokio::spawn({
1554            let metrics = metrics.clone();
1555            async move {
1556                loop {
1557                    if let Ok((mut socket, _)) = listener.accept().await {
1558                        let metrics = metrics.clone();
1559                        tokio::spawn(async move {
1560                            use tokio::io::{AsyncReadExt, AsyncWriteExt};
1561                            let mut buffer = [0; 1024];
1562                            let _ = socket.read(&mut buffer).await;
1563
1564                            let response = format!(
1565                                "HTTP/1.1 200 OK\r\nContent-Length: {}\r\n\r\n{}",
1566                                metrics.len(),
1567                                metrics
1568                            );
1569                            let _ = socket.write_all(response.as_bytes()).await;
1570                        });
1571                    }
1572                }
1573            }
1574        });
1575
1576        // Create a NetworkNode with the mock prometheus URI
1577        let mock_provider = Arc::new(MockNode::new());
1578        let mock_node = NetworkNode::new(
1579            "test_node",
1580            "ws://localhost:9944",
1581            "internal_ws_uri",
1582            &format!("http://127.0.0.1:{}/metrics", addr.port()),
1583            "/ip4/127.0.0.1/tcp/30333",
1584            NodeSpec::default(),
1585            mock_provider,
1586            GenCmdOptions::default(),
1587            NodeContext::Rc,
1588        );
1589
1590        // Get buckets with label filter
1591        let mut label_filters = HashMap::new();
1592        label_filters.insert("chain".to_string(), "rococo_local_testnet".to_string());
1593        let buckets = mock_node
1594            .get_histogram_buckets("substrate_block_verification_time", Some(label_filters))
1595            .await?;
1596
1597        // Should get the rococo_local_testnet chain's buckets
1598        assert_eq!(buckets.get("0.1"), Some(&10));
1599        assert_eq!(buckets.get("0.5"), Some(&15)); // 25 - 10
1600        assert_eq!(buckets.get("1.0"), Some(&10)); // 35 - 25
1601        assert_eq!(buckets.get("2.5"), Some(&5)); // 40 - 35
1602        assert_eq!(buckets.get("+Inf"), Some(&2)); // 42 - 40
1603
1604        // Get buckets with label filter for rococo
1605        let mut label_filters = std::collections::HashMap::new();
1606        label_filters.insert("chain".to_string(), "rococo_local_testnet".to_string());
1607
1608        let buckets_filtered = mock_node
1609            .get_histogram_buckets("substrate_block_verification_time", Some(label_filters))
1610            .await?;
1611
1612        assert_eq!(buckets_filtered.get("0.1"), Some(&10));
1613        assert_eq!(buckets_filtered.get("0.5"), Some(&15));
1614
1615        // Test 3: Get buckets with _bucket suffix already present
1616        let buckets_with_suffix = mock_node
1617            .get_histogram_buckets("substrate_block_verification_time_bucket", None)
1618            .await?;
1619
1620        assert_eq!(buckets_with_suffix.get("0.1"), Some(&10));
1621
1622        Ok(())
1623    }
1624
1625    #[tokio::test]
1626    async fn test_get_histogram_buckets_unordered() -> Result<(), anyhow::Error> {
1627        // Test that buckets are correctly sorted even when received out of order
1628        use std::sync::Arc;
1629
1630        let mock_metrics = concat!(
1631            "# HELP test_metric A test metric\n",
1632            "# TYPE test_metric histogram\n",
1633            "test_metric_bucket{le=\"2.5\"} 40\n",
1634            "test_metric_bucket{le=\"0.1\"} 10\n",
1635            "test_metric_bucket{le=\"+Inf\"} 42\n",
1636            "test_metric_bucket{le=\"1.0\"} 35\n",
1637            "test_metric_bucket{le=\"0.5\"} 25\n",
1638        );
1639
1640        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?;
1641        let addr = listener.local_addr()?;
1642        let metrics = Arc::new(mock_metrics.to_string());
1643
1644        tokio::spawn({
1645            let metrics = metrics.clone();
1646            async move {
1647                loop {
1648                    if let Ok((mut socket, _)) = listener.accept().await {
1649                        let metrics = metrics.clone();
1650                        tokio::spawn(async move {
1651                            use tokio::io::{AsyncReadExt, AsyncWriteExt};
1652                            let mut buffer = [0; 1024];
1653                            let _ = socket.read(&mut buffer).await;
1654                            let response = format!(
1655                                "HTTP/1.1 200 OK\r\nContent-Length: {}\r\n\r\n{}",
1656                                metrics.len(),
1657                                metrics
1658                            );
1659                            let _ = socket.write_all(response.as_bytes()).await;
1660                        });
1661                    }
1662                }
1663            }
1664        });
1665
1666        let mock_provider = Arc::new(MockNode::new());
1667        let mock_node = NetworkNode::new(
1668            "test_node",
1669            "ws://localhost:9944",
1670            "internal_ws_uri",
1671            &format!("http://127.0.0.1:{}/metrics", addr.port()),
1672            "/ip4/127.0.0.1/tcp/30333",
1673            NodeSpec::default(),
1674            mock_provider,
1675            GenCmdOptions::default(),
1676            NodeContext::Rc,
1677        );
1678
1679        let buckets = mock_node.get_histogram_buckets("test_metric", None).await?;
1680
1681        // Verify deltas are calculated correctly after sorting
1682        assert_eq!(buckets.get("0.1"), Some(&10)); // 10 - 0
1683        assert_eq!(buckets.get("0.5"), Some(&15)); // 25 - 10
1684        assert_eq!(buckets.get("1.0"), Some(&10)); // 35 - 25
1685        assert_eq!(buckets.get("2.5"), Some(&5)); // 40 - 35
1686        assert_eq!(buckets.get("+Inf"), Some(&2)); // 42 - 40
1687
1688        Ok(())
1689    }
1690
1691    #[tokio::test]
1692    async fn test_get_histogram_buckets_complex_labels() -> Result<(), anyhow::Error> {
1693        // Test label parsing with commas and special characters in values
1694        use std::sync::Arc;
1695
1696        let mock_metrics = concat!(
1697            "# HELP test_metric A test metric\n",
1698            "# TYPE test_metric histogram\n",
1699            "test_metric_bucket{method=\"GET,POST\",path=\"/api/test\",le=\"0.1\"} 5\n",
1700            "test_metric_bucket{method=\"GET,POST\",path=\"/api/test\",le=\"0.5\"} 15\n",
1701            "test_metric_bucket{method=\"GET,POST\",path=\"/api/test\",le=\"+Inf\"} 20\n",
1702        );
1703
1704        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?;
1705        let addr = listener.local_addr()?;
1706        let metrics = Arc::new(mock_metrics.to_string());
1707
1708        tokio::spawn({
1709            let metrics = metrics.clone();
1710            async move {
1711                loop {
1712                    if let Ok((mut socket, _)) = listener.accept().await {
1713                        let metrics = metrics.clone();
1714                        tokio::spawn(async move {
1715                            use tokio::io::{AsyncReadExt, AsyncWriteExt};
1716                            let mut buffer = [0; 1024];
1717                            let _ = socket.read(&mut buffer).await;
1718                            let response = format!(
1719                                "HTTP/1.1 200 OK\r\nContent-Length: {}\r\n\r\n{}",
1720                                metrics.len(),
1721                                metrics
1722                            );
1723                            let _ = socket.write_all(response.as_bytes()).await;
1724                        });
1725                    }
1726                }
1727            }
1728        });
1729
1730        let mock_provider = Arc::new(MockNode::new());
1731        let mock_node = NetworkNode::new(
1732            "test_node",
1733            "ws://localhost:9944",
1734            "internal_ws_uri",
1735            &format!("http://127.0.0.1:{}/metrics", addr.port()),
1736            "/ip4/127.0.0.1/tcp/30333",
1737            NodeSpec::default(),
1738            mock_provider,
1739            GenCmdOptions::default(),
1740            NodeContext::Rc,
1741        );
1742
1743        // Test without filter
1744        let buckets = mock_node.get_histogram_buckets("test_metric", None).await?;
1745        assert_eq!(buckets.get("0.1"), Some(&5));
1746        assert_eq!(buckets.get("0.5"), Some(&10)); // 15 - 5
1747        assert_eq!(buckets.get("+Inf"), Some(&5)); // 20 - 15
1748
1749        // Test with filter containing comma in value
1750        let mut label_filters = std::collections::HashMap::new();
1751        label_filters.insert("method".to_string(), "GET,POST".to_string());
1752
1753        let buckets_filtered = mock_node
1754            .get_histogram_buckets("test_metric", Some(label_filters))
1755            .await?;
1756
1757        assert_eq!(buckets_filtered.get("0.1"), Some(&5));
1758        assert_eq!(buckets_filtered.get("0.5"), Some(&10));
1759
1760        Ok(())
1761    }
1762
1763    #[test]
1764    fn test_compare_le_values() {
1765        use std::cmp::Ordering;
1766
1767        use crate::network::node::NetworkNode;
1768
1769        // Numeric comparison
1770        assert_eq!(NetworkNode::compare_le_values("0.1", "0.5"), Ordering::Less);
1771        assert_eq!(
1772            NetworkNode::compare_le_values("1.0", "0.5"),
1773            Ordering::Greater
1774        );
1775        assert_eq!(
1776            NetworkNode::compare_le_values("1.0", "1.0"),
1777            Ordering::Equal
1778        );
1779
1780        // +Inf handling
1781        assert_eq!(
1782            NetworkNode::compare_le_values("+Inf", "999"),
1783            Ordering::Greater
1784        );
1785        assert_eq!(
1786            NetworkNode::compare_le_values("0.1", "+Inf"),
1787            Ordering::Less
1788        );
1789        assert_eq!(
1790            NetworkNode::compare_le_values("+Inf", "+Inf"),
1791            Ordering::Equal
1792        );
1793
1794        // Large numbers
1795        assert_eq!(NetworkNode::compare_le_values("10", "100"), Ordering::Less);
1796        assert_eq!(
1797            NetworkNode::compare_le_values("1000", "999"),
1798            Ordering::Greater
1799        );
1800    }
1801
1802    fn mock_network_node(name: &str) -> NetworkNode {
1803        NetworkNode::new(
1804            name,
1805            "ws://127.0.0.1:9944",
1806            "ws://127.0.0.1:9944",
1807            "http://127.0.0.1:9615/metrics",
1808            "/ip4/127.0.0.1/tcp/30333",
1809            NodeSpec::default(),
1810            Arc::new(MockNode::new()),
1811            GenCmdOptions::default(),
1812            NodeContext::Rc,
1813        )
1814    }
1815
1816    fn mock_jam_node(name: &str) -> JamNetworkNode {
1817        JamNetworkNode::new(
1818            name,
1819            Arc::new(MockNode::new()),
1820            crate::network_spec::jamnode::JamNodeSpec::default(),
1821            std::net::IpAddr::from([127, 0, 0, 1]),
1822            GenCmdOptions::default(),
1823        )
1824    }
1825
1826    /// `NodeCore` is flattened into the node, so `zombie.json` keeps the same
1827    /// flat shape the attach path (`RawNetworkNode`) and the `{{node.field}}`
1828    /// replacements expect.
1829    #[test]
1830    fn test_network_node_serializes_flat() -> Result<(), anyhow::Error> {
1831        let value = serde_json::to_value(mock_network_node("alice"))?;
1832
1833        let obj = value.as_object().expect("node serializes as a map");
1834        assert_eq!(obj["name"], "alice");
1835        assert_eq!(obj["ws_uri"], "ws://127.0.0.1:9944");
1836        assert_eq!(obj["internal_ws_uri"], "ws://127.0.0.1:9944");
1837        assert_eq!(obj["kind"], "substrate");
1838        assert!(obj.contains_key("inner"), "inner should be flattened in");
1839
1840        // and it is still readable back by the attach path
1841        let raw: RawNetworkNode = serde_json::from_value(value)?;
1842        assert_eq!(raw.name, "alice");
1843        assert_eq!(raw.ws_uri, "ws://127.0.0.1:9944");
1844
1845        Ok(())
1846    }
1847
1848    #[test]
1849    fn test_jam_node_serializes_flat() -> Result<(), anyhow::Error> {
1850        let value = serde_json::to_value(mock_jam_node("jam-1"))?;
1851
1852        let obj = value.as_object().expect("node serializes as a map");
1853        assert_eq!(obj["name"], "jam-1");
1854        assert_eq!(obj["kind"], "jam");
1855
1856        let raw: jam::RawJamNetworkNode = serde_json::from_value(value)?;
1857        assert_eq!(raw.name, "jam-1");
1858
1859        Ok(())
1860    }
1861
1862    /// Both node types live in one registry and can be downcast back to their
1863    /// concrete type (trait upcasting to `dyn Any`, Rust >= 1.86).
1864    #[test]
1865    fn test_registry_holds_both_kinds_and_downcasts() {
1866        let mut registry: HashMap<String, Arc<dyn SpawnedNode>> = HashMap::new();
1867        registry.insert("alice".into(), Arc::new(mock_network_node("alice")));
1868        registry.insert("jam-1".into(), Arc::new(mock_jam_node("jam-1")));
1869
1870        let downcast = |name: &str| -> &dyn std::any::Any { registry[name].as_ref() };
1871
1872        // right kind -> concrete type back, with its own surface available
1873        let alice = downcast("alice")
1874            .downcast_ref::<NetworkNode>()
1875            .expect("alice is a substrate node");
1876        assert_eq!(alice.ws_uri(), "ws://127.0.0.1:9944");
1877
1878        let jam = downcast("jam-1")
1879            .downcast_ref::<JamNetworkNode>()
1880            .expect("jam-1 is a jam node");
1881        assert_eq!(jam.peer_addr(), "@127.0.0.1:0");
1882
1883        // wrong kind -> no downcast, and `kind()` says why
1884        assert!(downcast("alice").downcast_ref::<JamNetworkNode>().is_none());
1885        assert!(downcast("jam-1").downcast_ref::<NetworkNode>().is_none());
1886        assert_eq!(registry["alice"].kind(), NodeKind::Substrate);
1887        assert_eq!(registry["jam-1"].kind(), NodeKind::Jam);
1888
1889        // shared behaviour works through the erased type
1890        assert_eq!(registry["jam-1"].name(), "jam-1");
1891        assert!(!registry["jam-1"].is_running());
1892    }
1893}
zombienet_orchestrator/network/node.rs

zombienet_orchestrator/network/
node.rs
