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

zombienet_orchestrator/network/
node.rs
