Skip to main content

linera_service/cli_wrappers/
wallet.rs

1// Copyright (c) Zefchain Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use std::{
5    borrow::Cow,
6    collections::BTreeMap,
7    env,
8    marker::PhantomData,
9    mem,
10    path::{Path, PathBuf},
11    pin::Pin,
12    process::Stdio,
13    str::FromStr,
14    sync,
15    time::Duration,
16};
17
18use anyhow::{bail, ensure, Context, Result};
19use async_graphql::InputType;
20use async_tungstenite::tungstenite::{client::IntoClientRequest as _, http::HeaderValue};
21use futures::{SinkExt as _, Stream, StreamExt as _, TryStreamExt as _};
22use heck::ToKebabCase;
23use linera_base::{
24    abi::ContractAbi,
25    command::{resolve_binary, CommandExt},
26    crypto::{CryptoHash, InMemorySigner},
27    data_types::{Amount, BlockHeight, Bytecode, Epoch},
28    identifiers::{
29        Account, AccountOwner, ApplicationId, ChainId, IndexAndEvent, ModuleId, StreamId,
30    },
31    vm::VmRuntime,
32};
33use linera_client::client_options::ResourceControlPolicyConfig;
34use linera_core::worker::Notification;
35use linera_execution::committee::Committee;
36use linera_faucet_client::Faucet;
37use serde::{de::DeserializeOwned, ser::Serialize};
38use serde_command_opts::to_args;
39use serde_json::{json, Value};
40use tempfile::TempDir;
41use tokio::{
42    io::{AsyncBufReadExt, BufReader},
43    process::{Child, Command},
44    sync::oneshot,
45    task::JoinHandle,
46};
47use tracing::{error, info, warn};
48#[cfg(with_testing)]
49use {
50    futures::FutureExt as _,
51    linera_core::worker::Reason,
52    std::{collections::BTreeSet, future::Future},
53};
54
55use crate::{
56    cli::command::BenchmarkCommand,
57    cli_wrappers::{
58        local_net::{PathProvider, ProcessInbox},
59        Network,
60    },
61    util::{self, ChildExt},
62    Wallet,
63};
64
65/// The name of the environment variable that allows specifying additional arguments to be passed
66/// to the node-service command of the client.
67const CLIENT_SERVICE_ENV: &str = "LINERA_CLIENT_SERVICE_PARAMS";
68
69fn reqwest_client() -> reqwest::Client {
70    reqwest::ClientBuilder::new()
71        .timeout(Duration::from_secs(30))
72        .build()
73        .unwrap()
74}
75
76/// Wrapper to run a Linera client command.
77pub struct ClientWrapper {
78    binary_path: sync::Mutex<Option<PathBuf>>,
79    testing_prng_seed: Option<u64>,
80    storage: String,
81    wallet: String,
82    keystore: String,
83    network: Network,
84    /// Provides the working directory for the client and its files.
85    pub path_provider: PathProvider,
86    on_drop: OnClientDrop,
87    extra_args: Vec<String>,
88}
89
90/// Action to perform when the [`ClientWrapper`] is dropped.
91#[derive(Clone, Copy, Debug, Eq, PartialEq)]
92pub enum OnClientDrop {
93    /// Close all the chains on the wallet.
94    CloseChains,
95    /// Do not close any chains, leaving them active.
96    LeakChains,
97}
98
99impl ClientWrapper {
100    /// Creates a new [`ClientWrapper`], waiting for outgoing messages by default.
101    pub fn new(
102        path_provider: PathProvider,
103        network: Network,
104        testing_prng_seed: Option<u64>,
105        id: usize,
106        on_drop: OnClientDrop,
107    ) -> Self {
108        Self::new_with_extra_args(
109            path_provider,
110            network,
111            testing_prng_seed,
112            id,
113            on_drop,
114            vec!["--wait-for-outgoing-messages".to_string()],
115        )
116    }
117
118    /// Creates a new [`ClientWrapper`] with the given extra arguments passed to every command.
119    pub fn new_with_extra_args(
120        path_provider: PathProvider,
121        network: Network,
122        testing_prng_seed: Option<u64>,
123        id: usize,
124        on_drop: OnClientDrop,
125        extra_args: Vec<String>,
126    ) -> Self {
127        let storage = format!("rocksdb:{}/client_{id}.db", path_provider.path().display(),);
128        let wallet = format!("wallet_{id}.json");
129        let keystore = format!("keystore_{id}.json");
130        Self {
131            binary_path: sync::Mutex::new(None),
132            testing_prng_seed,
133            storage,
134            wallet,
135            keystore,
136            network,
137            path_provider,
138            on_drop,
139            extra_args,
140        }
141    }
142
143    /// Runs `linera project new`.
144    pub async fn project_new(&self, project_name: &str, linera_root: &Path) -> Result<TempDir> {
145        let tmp = TempDir::new()?;
146        let mut command = self.command().await?;
147        command
148            .current_dir(tmp.path())
149            .arg("project")
150            .arg("new")
151            .arg(project_name)
152            .arg("--linera-root")
153            .arg(linera_root)
154            .spawn_and_wait_for_stdout()
155            .await?;
156        Ok(tmp)
157    }
158
159    /// Runs `linera project publish`.
160    pub async fn project_publish<T: Serialize>(
161        &self,
162        path: PathBuf,
163        required_application_ids: Vec<String>,
164        publisher: impl Into<Option<ChainId>>,
165        argument: &T,
166    ) -> Result<String> {
167        let json_parameters = serde_json::to_string(&())?;
168        let json_argument = serde_json::to_string(argument)?;
169        let mut command = self.command().await?;
170        command
171            .arg("project")
172            .arg("publish-and-create")
173            .arg(path)
174            .args(publisher.into().iter().map(ChainId::to_string))
175            .args(["--json-parameters", &json_parameters])
176            .args(["--json-argument", &json_argument]);
177        if !required_application_ids.is_empty() {
178            command.arg("--required-application-ids");
179            command.args(required_application_ids);
180        }
181        let stdout = command.spawn_and_wait_for_stdout().await?;
182        Ok(stdout.trim().to_string())
183    }
184
185    /// Runs `linera project test`.
186    pub async fn project_test(&self, path: &Path) -> Result<()> {
187        self.command()
188            .await
189            .context("failed to create project test command")?
190            .current_dir(path)
191            .arg("project")
192            .arg("test")
193            .spawn_and_wait_for_stdout()
194            .await?;
195        Ok(())
196    }
197
198    async fn command_with_envs_and_arguments(
199        &self,
200        envs: &[(&str, &str)],
201        arguments: impl IntoIterator<Item = Cow<'_, str>>,
202    ) -> Result<Command> {
203        let mut command = self.command_binary().await?;
204        command.current_dir(self.path_provider.path());
205        for (key, value) in envs {
206            command.env(key, value);
207        }
208        for argument in arguments {
209            command.arg(&*argument);
210        }
211        Ok(command)
212    }
213
214    async fn command_with_envs(&self, envs: &[(&str, &str)]) -> Result<Command> {
215        self.command_with_envs_and_arguments(envs, self.command_arguments())
216            .await
217    }
218
219    async fn command_with_arguments(
220        &self,
221        arguments: impl IntoIterator<Item = Cow<'_, str>>,
222    ) -> Result<Command> {
223        self.command_with_envs_and_arguments(
224            &[(
225                "RUST_LOG",
226                &std::env::var("RUST_LOG").unwrap_or_else(|_| String::from("linera=debug")),
227            )],
228            arguments,
229        )
230        .await
231    }
232
233    async fn command(&self) -> Result<Command> {
234        self.command_with_envs(&[(
235            "RUST_LOG",
236            &std::env::var("RUST_LOG").unwrap_or_else(|_| String::from("linera=debug")),
237        )])
238        .await
239    }
240
241    fn required_command_arguments(&self) -> impl Iterator<Item = Cow<'_, str>> + '_ {
242        [
243            "--wallet".into(),
244            self.wallet.as_str().into(),
245            "--keystore".into(),
246            self.keystore.as_str().into(),
247            "--storage".into(),
248            self.storage.as_str().into(),
249            "--send-timeout-ms".into(),
250            "500000".into(),
251            "--recv-timeout-ms".into(),
252            "500000".into(),
253        ]
254        .into_iter()
255        .chain(self.extra_args.iter().map(|s| s.as_str().into()))
256    }
257
258    /// Returns an iterator over the arguments that should be added to all command invocations.
259    fn command_arguments(&self) -> impl Iterator<Item = Cow<'_, str>> + '_ {
260        self.required_command_arguments()
261            .chain(["--with-application-logs".into()])
262    }
263
264    /// Returns the [`Command`] instance configured to run the appropriate binary.
265    ///
266    /// The path is resolved once and cached inside `self` for subsequent usages.
267    async fn command_binary(&self) -> Result<Command> {
268        match self.command_with_cached_binary_path() {
269            Some(command) => Ok(command),
270            None => {
271                let resolved_path = resolve_binary("linera", env!("CARGO_PKG_NAME")).await?;
272                let command = Command::new(&resolved_path);
273
274                self.set_cached_binary_path(resolved_path);
275
276                Ok(command)
277            }
278        }
279    }
280
281    /// Returns a [`Command`] instance configured with the cached `binary_path`, if available.
282    fn command_with_cached_binary_path(&self) -> Option<Command> {
283        let binary_path = self.binary_path.lock().unwrap();
284
285        binary_path.as_ref().map(Command::new)
286    }
287
288    /// Sets the cached `binary_path` with the `new_binary_path`.
289    ///
290    /// # Panics
291    ///
292    /// If the cache is already set to a different value. In theory the two threads calling
293    /// `command_binary` can race and resolve the binary path twice, but they should always be the
294    /// same path.
295    fn set_cached_binary_path(&self, new_binary_path: PathBuf) {
296        let mut binary_path = self.binary_path.lock().unwrap();
297
298        if binary_path.is_none() {
299            *binary_path = Some(new_binary_path);
300        } else {
301            assert_eq!(*binary_path, Some(new_binary_path));
302        }
303    }
304
305    /// Runs `linera create-genesis-config`.
306    pub async fn create_genesis_config(
307        &self,
308        num_other_initial_chains: u32,
309        initial_funding: Amount,
310        policy_config: ResourceControlPolicyConfig,
311        http_allow_list: Option<Vec<String>>,
312    ) -> Result<()> {
313        let mut command = self.command().await?;
314        command
315            .args([
316                "create-genesis-config",
317                &num_other_initial_chains.to_string(),
318            ])
319            .args(["--initial-funding", &initial_funding.to_string()])
320            .args(["--committee", "committee.json"])
321            .args(["--genesis", "genesis.json"])
322            .args([
323                "--policy-config",
324                &policy_config.to_string().to_kebab_case(),
325            ]);
326        if let Some(allow_list) = http_allow_list {
327            command
328                .arg("--http-request-allow-list")
329                .arg(allow_list.join(","));
330        }
331        if let Some(seed) = self.testing_prng_seed {
332            command.arg("--testing-prng-seed").arg(seed.to_string());
333        }
334        command.spawn_and_wait_for_stdout().await?;
335        Ok(())
336    }
337
338    /// Runs `linera wallet init`. The genesis config is read from `genesis.json`, or from the
339    /// faucet if provided.
340    pub async fn wallet_init(&self, faucet: Option<&'_ Faucet>) -> Result<()> {
341        let mut command = self.command().await?;
342        command.args(["wallet", "init"]);
343        match faucet {
344            None => command.args(["--genesis", "genesis.json"]),
345            Some(faucet) => command.args(["--faucet", faucet.url()]),
346        };
347        if let Some(seed) = self.testing_prng_seed {
348            command.arg("--testing-prng-seed").arg(seed.to_string());
349        }
350        command.spawn_and_wait_for_stdout().await?;
351        Ok(())
352    }
353
354    /// Runs `linera wallet request-chain`.
355    pub async fn request_chain(
356        &self,
357        faucet: &Faucet,
358        set_default: bool,
359    ) -> Result<(ChainId, AccountOwner)> {
360        let mut command = self.command().await?;
361        command.args(["wallet", "request-chain", "--faucet", faucet.url()]);
362        if set_default {
363            command.arg("--set-default");
364        }
365        let stdout = command.spawn_and_wait_for_stdout().await?;
366        let mut lines = stdout.split_whitespace();
367        let chain_id: ChainId = lines.next().context("missing chain ID")?.parse()?;
368        let owner = lines.next().context("missing chain owner")?.parse()?;
369        Ok((chain_id, owner))
370    }
371
372    /// Runs `linera wallet publish-and-create`.
373    #[expect(clippy::too_many_arguments)]
374    pub async fn publish_and_create<
375        A: ContractAbi,
376        Parameters: Serialize,
377        InstantiationArgument: Serialize,
378    >(
379        &self,
380        contract: PathBuf,
381        service: PathBuf,
382        vm_runtime: VmRuntime,
383        parameters: &Parameters,
384        argument: &InstantiationArgument,
385        required_application_ids: &[ApplicationId],
386        publisher: impl Into<Option<ChainId>>,
387    ) -> Result<ApplicationId<A>> {
388        let json_parameters = serde_json::to_string(parameters)?;
389        let json_argument = serde_json::to_string(argument)?;
390        let mut command = self.command().await?;
391        let vm_runtime = format!("{vm_runtime}");
392        command
393            .arg("publish-and-create")
394            .args([contract, service])
395            .args(["--vm-runtime", &vm_runtime.to_lowercase()])
396            .args(publisher.into().iter().map(ChainId::to_string))
397            .args(["--json-parameters", &json_parameters])
398            .args(["--json-argument", &json_argument]);
399        if !required_application_ids.is_empty() {
400            command.arg("--required-application-ids");
401            command.args(
402                required_application_ids
403                    .iter()
404                    .map(ApplicationId::to_string),
405            );
406        }
407        let stdout = command.spawn_and_wait_for_stdout().await?;
408        Ok(stdout.trim().parse::<ApplicationId>()?.with_abi())
409    }
410
411    /// Runs `linera publish-module`.
412    pub async fn publish_module<Abi, Parameters, InstantiationArgument>(
413        &self,
414        contract: PathBuf,
415        service: PathBuf,
416        vm_runtime: VmRuntime,
417        publisher: impl Into<Option<ChainId>>,
418    ) -> Result<ModuleId<Abi, Parameters, InstantiationArgument>> {
419        let stdout = self
420            .command()
421            .await?
422            .arg("publish-module")
423            .args([contract, service])
424            .args(["--vm-runtime", &format!("{vm_runtime}").to_lowercase()])
425            .args(publisher.into().iter().map(ChainId::to_string))
426            .spawn_and_wait_for_stdout()
427            .await?;
428        let module_id: ModuleId = stdout.trim().parse()?;
429        Ok(module_id.with_abi())
430    }
431
432    /// Runs `linera publish-module-with-formats`.
433    #[allow(clippy::too_many_arguments)]
434    pub async fn publish_module_with_formats<Abi, Parameters, InstantiationArgument>(
435        &self,
436        contract: PathBuf,
437        service: PathBuf,
438        formats: PathBuf,
439        registry_application_id: ApplicationId,
440        vm_runtime: VmRuntime,
441        publisher: impl Into<Option<ChainId>>,
442    ) -> Result<ModuleId<Abi, Parameters, InstantiationArgument>> {
443        let stdout = self
444            .command()
445            .await?
446            .arg("publish-module-with-formats")
447            .args([contract, service, formats])
448            .arg(registry_application_id.to_string())
449            .args(["--vm-runtime", &format!("{vm_runtime}").to_lowercase()])
450            .args(publisher.into().iter().map(ChainId::to_string))
451            .spawn_and_wait_for_stdout()
452            .await?;
453        let module_id: ModuleId = stdout.trim().parse()?;
454        Ok(module_id.with_abi())
455    }
456
457    /// Runs `linera create-application`.
458    pub async fn create_application<
459        Abi: ContractAbi,
460        Parameters: Serialize,
461        InstantiationArgument: Serialize,
462    >(
463        &self,
464        module_id: &ModuleId<Abi, Parameters, InstantiationArgument>,
465        parameters: &Parameters,
466        argument: &InstantiationArgument,
467        required_application_ids: &[ApplicationId],
468        creator: impl Into<Option<ChainId>>,
469    ) -> Result<ApplicationId<Abi>> {
470        let json_parameters = serde_json::to_string(parameters)?;
471        let json_argument = serde_json::to_string(argument)?;
472        let mut command = self.command().await?;
473        command
474            .arg("create-application")
475            .arg(module_id.forget_abi().to_string())
476            .args(["--json-parameters", &json_parameters])
477            .args(["--json-argument", &json_argument])
478            .args(creator.into().iter().map(ChainId::to_string));
479        if !required_application_ids.is_empty() {
480            command.arg("--required-application-ids");
481            command.args(
482                required_application_ids
483                    .iter()
484                    .map(ApplicationId::to_string),
485            );
486        }
487        let stdout = command.spawn_and_wait_for_stdout().await?;
488        Ok(stdout.trim().parse::<ApplicationId>()?.with_abi())
489    }
490
491    /// Runs `linera service`.
492    pub async fn run_node_service(
493        &self,
494        port: impl Into<Option<u16>>,
495        process_inbox: ProcessInbox,
496    ) -> Result<NodeService> {
497        self.run_node_service_with_options(port, process_inbox, &[], &[], false)
498            .await
499    }
500
501    /// Runs `linera service` with optional task processor configuration.
502    pub async fn run_node_service_with_options(
503        &self,
504        port: impl Into<Option<u16>>,
505        process_inbox: ProcessInbox,
506        operator_application_ids: &[ApplicationId],
507        operators: &[(String, PathBuf)],
508        read_only: bool,
509    ) -> Result<NodeService> {
510        self.run_node_service_with_all_options(
511            port,
512            process_inbox,
513            operator_application_ids,
514            operators,
515            read_only,
516            &[],
517            &[],
518        )
519        .await
520    }
521
522    /// Runs `linera service` with all available options.
523    #[expect(clippy::too_many_arguments)]
524    pub async fn run_node_service_with_all_options(
525        &self,
526        port: impl Into<Option<u16>>,
527        process_inbox: ProcessInbox,
528        operator_application_ids: &[ApplicationId],
529        operators: &[(String, PathBuf)],
530        read_only: bool,
531        allowed_subscriptions: &[String],
532        subscription_ttls: &[(String, u64)],
533    ) -> Result<NodeService> {
534        let port = port.into().unwrap_or(8080);
535        let mut command = self.command().await?;
536        command.arg("service");
537        if let ProcessInbox::Skip = process_inbox {
538            command.arg("--listener-skip-process-inbox");
539        }
540        if let Ok(var) = env::var(CLIENT_SERVICE_ENV) {
541            command.args(var.split_whitespace());
542        }
543        for app_id in operator_application_ids {
544            command.args(["--operator-application-ids", &app_id.to_string()]);
545        }
546        for (name, path) in operators {
547            command.args(["--operators", &format!("{}={}", name, path.display())]);
548        }
549        if read_only {
550            command.arg("--read-only");
551        }
552        for query in allowed_subscriptions {
553            command.args(["--allow-subscription", query]);
554        }
555        for (name, secs) in subscription_ttls {
556            command.args(["--subscription-ttl-secs", &format!("{name}={secs}")]);
557        }
558        let child = command
559            .args(["--port".to_string(), port.to_string()])
560            .spawn_into()?;
561        let client = reqwest_client();
562        for i in 0..10 {
563            linera_base::time::timer::sleep(Duration::from_secs(i)).await;
564            let request = client.get(format!("http://localhost:{port}/")).send().await;
565            if request.is_ok() {
566                info!("Node service has started");
567                return Ok(NodeService::new(port, child));
568            } else {
569                warn!("Waiting for node service to start");
570            }
571        }
572        bail!("Failed to start node service");
573    }
574
575    /// Runs `linera service` with a controller application.
576    pub async fn run_node_service_with_controller(
577        &self,
578        port: impl Into<Option<u16>>,
579        process_inbox: ProcessInbox,
580        controller_id: &ApplicationId,
581        operators: &[(String, PathBuf)],
582    ) -> Result<NodeService> {
583        let port = port.into().unwrap_or(8080);
584        let mut command = self.command().await?;
585        command.arg("service");
586        if let ProcessInbox::Skip = process_inbox {
587            command.arg("--listener-skip-process-inbox");
588        }
589        if let Ok(var) = env::var(CLIENT_SERVICE_ENV) {
590            command.args(var.split_whitespace());
591        }
592        command.args(["--controller-id", &controller_id.to_string()]);
593        for (name, path) in operators {
594            command.args(["--operators", &format!("{}={}", name, path.display())]);
595        }
596        let child = command
597            .args(["--port".to_string(), port.to_string()])
598            .spawn_into()?;
599        let client = reqwest_client();
600        for i in 0..10 {
601            linera_base::time::timer::sleep(Duration::from_secs(i)).await;
602            let request = client.get(format!("http://localhost:{port}/")).send().await;
603            if request.is_ok() {
604                tracing::info!("Node service has started");
605                return Ok(NodeService::new(port, child));
606            } else {
607                tracing::warn!("Waiting for node service to start");
608            }
609        }
610        bail!("Failed to start node service");
611    }
612
613    /// Runs `linera validator query`
614    pub async fn query_validator(&self, address: &str) -> Result<CryptoHash> {
615        let mut command = self.command().await?;
616        command.arg("validator").arg("query").arg(address);
617        let stdout = command.spawn_and_wait_for_stdout().await?;
618
619        // Parse the genesis config hash from the output.
620        // It's on a line like "Genesis config hash: <hash>"
621        let hash = stdout
622            .lines()
623            .find_map(|line| {
624                line.strip_prefix("Genesis config hash: ")
625                    .and_then(|hash_str| hash_str.trim().parse().ok())
626            })
627            .context("error while parsing the result of `linera validator query`")?;
628        Ok(hash)
629    }
630
631    /// Runs `linera validator list`.
632    pub async fn query_validators(&self, chain_id: Option<ChainId>) -> Result<()> {
633        let mut command = self.command().await?;
634        command.arg("validator").arg("list");
635        if let Some(chain_id) = chain_id {
636            command.args(["--chain-id", &chain_id.to_string()]);
637        }
638        command.spawn_and_wait_for_stdout().await?;
639        Ok(())
640    }
641
642    /// Runs `linera validator sync`.
643    pub async fn sync_validator(
644        &self,
645        chain_ids: impl IntoIterator<Item = &ChainId>,
646        validator_address: impl Into<String>,
647    ) -> Result<()> {
648        let mut command = self.command().await?;
649        command
650            .arg("validator")
651            .arg("sync")
652            .arg(validator_address.into());
653        let mut chain_ids = chain_ids.into_iter().peekable();
654        if chain_ids.peek().is_some() {
655            command
656                .arg("--chains")
657                .args(chain_ids.map(ChainId::to_string));
658        }
659        command.spawn_and_wait_for_stdout().await?;
660        Ok(())
661    }
662
663    /// Runs `linera validator benchmark` and returns its stdout.
664    pub async fn validator_benchmark(
665        &self,
666        address: impl Into<String>,
667        chains: impl IntoIterator<Item = &ChainId>,
668        extra_args: &[&str],
669    ) -> Result<String> {
670        let mut command = self.command().await?;
671        command
672            .arg("validator")
673            .arg("benchmark")
674            .arg(address.into());
675        for chain in chains {
676            command.args(["--chain", &chain.to_string()]);
677        }
678        command.args(extra_args);
679        command.spawn_and_wait_for_stdout().await
680    }
681
682    /// Runs `linera faucet`.
683    pub async fn run_faucet(
684        &self,
685        port: impl Into<Option<u16>>,
686        chain_id: ChainId,
687        amount: Amount,
688    ) -> Result<FaucetService> {
689        let port = port.into().unwrap_or(8080);
690        let temp_dir = tempfile::tempdir()
691            .context("Failed to create temporary directory for faucet storage")?;
692        let storage_path = temp_dir.path().join("faucet_storage.sqlite");
693        let mut command = self.command().await?;
694        let child = command
695            .arg("faucet")
696            .arg(chain_id.to_string())
697            .args(["--port".to_string(), port.to_string()])
698            .args(["--amount".to_string(), amount.to_string()])
699            .args([
700                "--storage-path".to_string(),
701                storage_path.to_string_lossy().to_string(),
702            ])
703            .spawn_into()?;
704        let client = reqwest_client();
705        for i in 0..10 {
706            linera_base::time::timer::sleep(Duration::from_secs(i)).await;
707            let request = client.get(format!("http://localhost:{port}/")).send().await;
708            if request.is_ok() {
709                info!("Faucet has started");
710                return Ok(FaucetService::new(port, child, temp_dir));
711            } else {
712                warn!("Waiting for faucet to start");
713            }
714        }
715        bail!("Failed to start faucet");
716    }
717
718    /// Runs `linera local-balance`.
719    pub async fn local_balance(&self, account: Account) -> Result<Amount> {
720        let stdout = self
721            .command()
722            .await?
723            .arg("local-balance")
724            .arg(account.to_string())
725            .spawn_and_wait_for_stdout()
726            .await?;
727        let amount = stdout
728            .trim()
729            .parse()
730            .context("error while parsing the result of `linera local-balance`")?;
731        Ok(amount)
732    }
733
734    /// Runs `linera query-balance`.
735    pub async fn query_balance(&self, account: Account) -> Result<Amount> {
736        let stdout = self
737            .command()
738            .await?
739            .arg("query-balance")
740            .arg(account.to_string())
741            .spawn_and_wait_for_stdout()
742            .await?;
743        let amount = stdout
744            .trim()
745            .parse()
746            .context("error while parsing the result of `linera query-balance`")?;
747        Ok(amount)
748    }
749
750    /// Runs `linera sync`.
751    pub async fn sync(&self, chain_id: ChainId) -> Result<()> {
752        self.command()
753            .await?
754            .arg("sync")
755            .arg(chain_id.to_string())
756            .spawn_and_wait_for_stdout()
757            .await?;
758        Ok(())
759    }
760
761    /// Runs `linera process-inbox`.
762    pub async fn process_inbox(&self, chain_id: ChainId) -> Result<()> {
763        self.command()
764            .await?
765            .arg("process-inbox")
766            .arg(chain_id.to_string())
767            .spawn_and_wait_for_stdout()
768            .await?;
769        Ok(())
770    }
771
772    /// Runs `linera transfer`.
773    pub async fn transfer(&self, amount: Amount, from: ChainId, to: ChainId) -> Result<()> {
774        self.command()
775            .await?
776            .arg("transfer")
777            .arg(amount.to_string())
778            .args(["--from", &from.to_string()])
779            .args(["--to", &to.to_string()])
780            .spawn_and_wait_for_stdout()
781            .await?;
782        Ok(())
783    }
784
785    /// Runs `linera transfer` with no logging.
786    pub async fn transfer_with_silent_logs(
787        &self,
788        amount: Amount,
789        from: ChainId,
790        to: ChainId,
791    ) -> Result<()> {
792        self.command()
793            .await?
794            .env("RUST_LOG", "off")
795            .arg("transfer")
796            .arg(amount.to_string())
797            .args(["--from", &from.to_string()])
798            .args(["--to", &to.to_string()])
799            .spawn_and_wait_for_stdout()
800            .await?;
801        Ok(())
802    }
803
804    /// Runs `linera transfer` with owner accounts.
805    pub async fn transfer_with_accounts(
806        &self,
807        amount: Amount,
808        from: Account,
809        to: Account,
810    ) -> Result<()> {
811        self.command()
812            .await?
813            .arg("transfer")
814            .arg(amount.to_string())
815            .args(["--from", &from.to_string()])
816            .args(["--to", &to.to_string()])
817            .spawn_and_wait_for_stdout()
818            .await?;
819        Ok(())
820    }
821
822    fn benchmark_command_internal(command: &mut Command, args: &BenchmarkCommand) -> Result<()> {
823        let mut formatted_args = to_args(args)?;
824        let subcommand = formatted_args.remove(0);
825        // The subcommand is followed by the flattened options, which are preceded by "options".
826        // So remove that as well.
827        formatted_args.remove(0);
828        let options = formatted_args
829            .chunks_exact(2)
830            .flat_map(|pair| {
831                let option = format!("--{}", pair[0]);
832                match pair[1].as_str() {
833                    "true" => vec![option],
834                    "false" => vec![],
835                    _ => vec![option, pair[1].clone()],
836                }
837            })
838            .collect::<Vec<_>>();
839        command
840            .args([
841                "--max-pending-message-bundles",
842                &args.transactions_per_block().to_string(),
843            ])
844            .arg("benchmark")
845            .arg(subcommand)
846            .args(options);
847        Ok(())
848    }
849
850    async fn benchmark_command_with_envs(
851        &self,
852        args: BenchmarkCommand,
853        envs: &[(&str, &str)],
854    ) -> Result<Command> {
855        let mut command = self
856            .command_with_envs_and_arguments(envs, self.required_command_arguments())
857            .await?;
858        Self::benchmark_command_internal(&mut command, &args)?;
859        Ok(command)
860    }
861
862    async fn benchmark_command(&self, args: BenchmarkCommand) -> Result<Command> {
863        let mut command = self
864            .command_with_arguments(self.required_command_arguments())
865            .await?;
866        Self::benchmark_command_internal(&mut command, &args)?;
867        Ok(command)
868    }
869
870    /// Runs `linera benchmark`.
871    pub async fn benchmark(&self, args: BenchmarkCommand) -> Result<()> {
872        let mut command = self.benchmark_command(args).await?;
873        command.spawn_and_wait_for_stdout().await?;
874        Ok(())
875    }
876
877    /// Runs `linera benchmark`, but detached: don't wait for the command to finish, just spawn it
878    /// and return the child process, and the handles to the stdout and stderr.
879    pub async fn benchmark_detached(
880        &self,
881        args: BenchmarkCommand,
882        tx: oneshot::Sender<()>,
883    ) -> Result<(Child, JoinHandle<()>, JoinHandle<()>)> {
884        let mut child = self
885            .benchmark_command_with_envs(args, &[("RUST_LOG", "linera=info")])
886            .await?
887            .kill_on_drop(true)
888            .stdin(Stdio::piped())
889            .stdout(Stdio::piped())
890            .stderr(Stdio::piped())
891            .spawn()?;
892
893        let pid = child.id().expect("failed to get pid");
894        let stdout = child.stdout.take().expect("stdout not open");
895        let stdout_handle = tokio::spawn(async move {
896            let mut lines = BufReader::new(stdout).lines();
897            while let Ok(Some(line)) = lines.next_line().await {
898                println!("benchmark{{pid={pid}}} {line}");
899            }
900        });
901
902        let stderr = child.stderr.take().expect("stderr not open");
903        let stderr_handle = tokio::spawn(async move {
904            let mut lines = BufReader::new(stderr).lines();
905            let mut tx = Some(tx);
906            while let Ok(Some(line)) = lines.next_line().await {
907                if line.contains("Ready to start benchmark") {
908                    tx.take()
909                        .expect("Should only send signal once")
910                        .send(())
911                        .expect("failed to send ready signal to main thread");
912                } else {
913                    println!("benchmark{{pid={pid}}} {line}");
914                }
915            }
916        });
917        Ok((child, stdout_handle, stderr_handle))
918    }
919
920    async fn open_chain_internal(
921        &self,
922        from: ChainId,
923        owner: Option<AccountOwner>,
924        initial_balance: Amount,
925        super_owner: bool,
926    ) -> Result<(ChainId, AccountOwner)> {
927        let mut command = self.command().await?;
928        command
929            .arg("open-chain")
930            .args(["--from", &from.to_string()])
931            .args(["--initial-balance", &initial_balance.to_string()]);
932
933        if let Some(owner) = owner {
934            command.args(["--owner", &owner.to_string()]);
935        }
936
937        if super_owner {
938            command.arg("--super-owner");
939        }
940
941        let stdout = command.spawn_and_wait_for_stdout().await?;
942        let mut split = stdout.split('\n');
943        let chain_id = ChainId::from_str(split.next().context("no chain ID in output")?)?;
944        let new_owner = AccountOwner::from_str(split.next().context("no owner in output")?)?;
945        if let Some(owner) = owner {
946            assert_eq!(owner, new_owner);
947        }
948        Ok((chain_id, new_owner))
949    }
950
951    /// Runs `linera open-chain --super-owner`.
952    pub async fn open_chain_super_owner(
953        &self,
954        from: ChainId,
955        owner: Option<AccountOwner>,
956        initial_balance: Amount,
957    ) -> Result<(ChainId, AccountOwner)> {
958        self.open_chain_internal(from, owner, initial_balance, true)
959            .await
960    }
961
962    /// Runs `linera open-chain`.
963    pub async fn open_chain(
964        &self,
965        from: ChainId,
966        owner: Option<AccountOwner>,
967        initial_balance: Amount,
968    ) -> Result<(ChainId, AccountOwner)> {
969        self.open_chain_internal(from, owner, initial_balance, false)
970            .await
971    }
972
973    /// Runs `linera open-chain` then `linera assign`.
974    pub async fn open_and_assign(
975        &self,
976        client: &ClientWrapper,
977        initial_balance: Amount,
978    ) -> Result<ChainId> {
979        let our_chain = self
980            .load_wallet()?
981            .default_chain()
982            .context("no default chain found")?;
983        let owner = client.keygen().await?;
984        let (new_chain, _) = self
985            .open_chain(our_chain, Some(owner), initial_balance)
986            .await?;
987        client.assign(owner, new_chain).await?;
988        Ok(new_chain)
989    }
990
991    /// Runs `linera open-multi-owner-chain` and returns the new chain's ID.
992    pub async fn open_multi_owner_chain(
993        &self,
994        from: ChainId,
995        owners: BTreeMap<AccountOwner, u64>,
996        multi_leader_rounds: u32,
997        balance: Amount,
998        base_timeout_ms: u64,
999    ) -> Result<ChainId> {
1000        let mut command = self.command().await?;
1001        command
1002            .arg("open-multi-owner-chain")
1003            .args(["--from", &from.to_string()])
1004            .arg("--owners")
1005            .arg(serde_json::to_string(&owners)?)
1006            .args(["--base-timeout-ms", &base_timeout_ms.to_string()]);
1007        command
1008            .args(["--multi-leader-rounds", &multi_leader_rounds.to_string()])
1009            .args(["--initial-balance", &balance.to_string()]);
1010
1011        let stdout = command.spawn_and_wait_for_stdout().await?;
1012        let mut split = stdout.split('\n');
1013        let chain_id = ChainId::from_str(split.next().context("no chain ID in output")?)?;
1014
1015        Ok(chain_id)
1016    }
1017
1018    /// Runs `linera change-ownership`.
1019    pub async fn change_ownership(
1020        &self,
1021        chain_id: ChainId,
1022        super_owners: Vec<AccountOwner>,
1023        owners: Vec<AccountOwner>,
1024    ) -> Result<()> {
1025        let mut command = self.command().await?;
1026        command
1027            .arg("change-ownership")
1028            .args(["--chain-id", &chain_id.to_string()]);
1029        command
1030            .arg("--super-owners")
1031            .arg(serde_json::to_string(&super_owners)?);
1032        command.arg("--owners").arg(serde_json::to_string(
1033            &owners
1034                .into_iter()
1035                .zip(std::iter::repeat(100u64))
1036                .collect::<BTreeMap<_, _>>(),
1037        )?);
1038        command.spawn_and_wait_for_stdout().await?;
1039        Ok(())
1040    }
1041
1042    /// Runs `linera wallet follow-chain CHAIN_ID`.
1043    pub async fn follow_chain(&self, chain_id: ChainId, sync: bool) -> Result<()> {
1044        let mut command = self.command().await?;
1045        command
1046            .args(["wallet", "follow-chain"])
1047            .arg(chain_id.to_string());
1048        if sync {
1049            command.arg("--sync");
1050        }
1051        command.spawn_and_wait_for_stdout().await?;
1052        Ok(())
1053    }
1054
1055    /// Runs `linera wallet forget-chain CHAIN_ID`.
1056    pub async fn forget_chain(&self, chain_id: ChainId) -> Result<()> {
1057        let mut command = self.command().await?;
1058        command
1059            .args(["wallet", "forget-chain"])
1060            .arg(chain_id.to_string());
1061        command.spawn_and_wait_for_stdout().await?;
1062        Ok(())
1063    }
1064
1065    /// Runs `linera wallet set-default CHAIN_ID`.
1066    pub async fn set_default_chain(&self, chain_id: ChainId) -> Result<()> {
1067        let mut command = self.command().await?;
1068        command
1069            .args(["wallet", "set-default"])
1070            .arg(chain_id.to_string());
1071        command.spawn_and_wait_for_stdout().await?;
1072        Ok(())
1073    }
1074
1075    /// Runs `linera retry-pending-block` and returns the hash of the produced block, if any.
1076    pub async fn retry_pending_block(
1077        &self,
1078        chain_id: Option<ChainId>,
1079    ) -> Result<Option<CryptoHash>> {
1080        let mut command = self.command().await?;
1081        command.arg("retry-pending-block");
1082        if let Some(chain_id) = chain_id {
1083            command.arg(chain_id.to_string());
1084        }
1085        let stdout = command.spawn_and_wait_for_stdout().await?;
1086        let stdout = stdout.trim();
1087        if stdout.is_empty() {
1088            Ok(None)
1089        } else {
1090            Ok(Some(CryptoHash::from_str(stdout)?))
1091        }
1092    }
1093
1094    /// Runs `linera publish-data-blob`.
1095    pub async fn publish_data_blob(
1096        &self,
1097        path: &Path,
1098        chain_id: Option<ChainId>,
1099    ) -> Result<CryptoHash> {
1100        let mut command = self.command().await?;
1101        command.arg("publish-data-blob").arg(path);
1102        if let Some(chain_id) = chain_id {
1103            command.arg(chain_id.to_string());
1104        }
1105        let stdout = command.spawn_and_wait_for_stdout().await?;
1106        let stdout = stdout.trim();
1107        Ok(CryptoHash::from_str(stdout)?)
1108    }
1109
1110    /// Runs `linera read-data-blob`.
1111    pub async fn read_data_blob(&self, hash: CryptoHash, chain_id: Option<ChainId>) -> Result<()> {
1112        let mut command = self.command().await?;
1113        command.arg("read-data-blob").arg(hash.to_string());
1114        if let Some(chain_id) = chain_id {
1115            command.arg(chain_id.to_string());
1116        }
1117        command.spawn_and_wait_for_stdout().await?;
1118        Ok(())
1119    }
1120
1121    /// Loads the wallet from this client's wallet file.
1122    pub fn load_wallet(&self) -> Result<Wallet> {
1123        Ok(Wallet::read(&self.wallet_path())?)
1124    }
1125
1126    /// Loads the in-memory signer from this client's keystore file.
1127    pub fn load_keystore(&self) -> Result<InMemorySigner> {
1128        util::read_json(self.keystore_path())
1129    }
1130
1131    /// Returns the path to this client's wallet file.
1132    pub fn wallet_path(&self) -> PathBuf {
1133        self.path_provider.path().join(&self.wallet)
1134    }
1135
1136    /// Returns the path to this client's keystore file.
1137    pub fn keystore_path(&self) -> PathBuf {
1138        self.path_provider.path().join(&self.keystore)
1139    }
1140
1141    /// Returns the storage specification string used by this client.
1142    pub fn storage_path(&self) -> &str {
1143        &self.storage
1144    }
1145
1146    /// Returns the owner of the wallet's default chain, if any.
1147    pub fn get_owner(&self) -> Option<AccountOwner> {
1148        let wallet = self.load_wallet().ok()?;
1149        wallet
1150            .get(wallet.default_chain()?)
1151            .expect("default chain must be in wallet")
1152            .owner
1153    }
1154
1155    /// Returns whether the given chain is present in the wallet.
1156    pub fn is_chain_present_in_wallet(&self, chain: ChainId) -> bool {
1157        self.load_wallet()
1158            .ok()
1159            .is_some_and(|wallet| wallet.get(chain).is_some())
1160    }
1161
1162    /// Runs `linera validator add`.
1163    pub async fn set_validator(
1164        &self,
1165        validator_key: &(String, String),
1166        port: usize,
1167        votes: usize,
1168    ) -> Result<()> {
1169        let address = format!("{}:127.0.0.1:{}", self.network.short(), port);
1170        self.command()
1171            .await?
1172            .arg("validator")
1173            .arg("add")
1174            .args(["--public-key", &validator_key.0])
1175            .args(["--account-key", &validator_key.1])
1176            .args(["--address", &address])
1177            .args(["--votes", &votes.to_string()])
1178            .spawn_and_wait_for_stdout()
1179            .await?;
1180        Ok(())
1181    }
1182
1183    /// Runs `linera validator remove`.
1184    pub async fn remove_validator(&self, validator_key: &str) -> Result<()> {
1185        self.command()
1186            .await?
1187            .arg("validator")
1188            .arg("remove")
1189            .args(["--public-key", validator_key])
1190            .spawn_and_wait_for_stdout()
1191            .await?;
1192        Ok(())
1193    }
1194
1195    /// Runs `linera validator update` to add, modify, and remove validators in one call.
1196    pub async fn change_validators(
1197        &self,
1198        add_validators: &[(String, String, usize, usize)], // (public_key, account_key, port, votes)
1199        modify_validators: &[(String, String, usize, usize)], // (public_key, account_key, port, votes)
1200        remove_validators: &[String],
1201    ) -> Result<()> {
1202        use std::str::FromStr;
1203
1204        use linera_base::crypto::{AccountPublicKey, ValidatorPublicKey};
1205
1206        // Build a map that will be serialized to JSON
1207        // Use the exact types that deserialization expects
1208        let mut changes = std::collections::HashMap::new();
1209
1210        // Add/modify validators
1211        for (public_key_str, account_key_str, port, votes) in
1212            add_validators.iter().chain(modify_validators.iter())
1213        {
1214            let public_key = ValidatorPublicKey::from_str(public_key_str)
1215                .with_context(|| format!("Invalid validator public key: {public_key_str}"))?;
1216
1217            let account_key = AccountPublicKey::from_str(account_key_str)
1218                .with_context(|| format!("Invalid account public key: {account_key_str}"))?;
1219
1220            let address = format!("{}:127.0.0.1:{}", self.network.short(), port)
1221                .parse()
1222                .unwrap();
1223
1224            // Create ValidatorChange struct
1225            let change = crate::cli::validator::Change {
1226                account_key,
1227                address,
1228                votes: crate::cli::validator::Votes(
1229                    std::num::NonZero::new(*votes as u64).context("Votes must be non-zero")?,
1230                ),
1231            };
1232
1233            changes.insert(public_key, Some(change));
1234        }
1235
1236        // Remove validators (set to None)
1237        for validator_key_str in remove_validators {
1238            let public_key = ValidatorPublicKey::from_str(validator_key_str)
1239                .with_context(|| format!("Invalid validator public key: {validator_key_str}"))?;
1240            changes.insert(public_key, None);
1241        }
1242
1243        // Create temporary file with JSON
1244        let temp_file = tempfile::NamedTempFile::new()
1245            .context("Failed to create temporary file for validator changes")?;
1246        serde_json::to_writer(&temp_file, &changes)
1247            .context("Failed to write validator changes to file")?;
1248        let temp_path = temp_file.path();
1249
1250        self.command()
1251            .await?
1252            .arg("validator")
1253            .arg("update")
1254            .arg(temp_path)
1255            .arg("--yes") // Skip confirmation prompt
1256            .spawn_and_wait_for_stdout()
1257            .await?;
1258
1259        Ok(())
1260    }
1261
1262    /// Runs `linera revoke-epochs`.
1263    pub async fn revoke_epochs(&self, epoch: Epoch) -> Result<()> {
1264        self.command()
1265            .await?
1266            .arg("revoke-epochs")
1267            .arg(epoch.to_string())
1268            .spawn_and_wait_for_stdout()
1269            .await?;
1270        Ok(())
1271    }
1272
1273    /// Runs `linera resource-control-policy --http-request-allow-list`.
1274    pub async fn change_http_whitelist(&self, list: &[&str]) -> Result<()> {
1275        self.command()
1276            .await?
1277            .arg("resource-control-policy")
1278            .arg("--http-request-allow-list")
1279            .arg(list.join(","))
1280            .spawn_and_wait_for_stdout()
1281            .await?;
1282        Ok(())
1283    }
1284
1285    /// Runs `linera keygen`.
1286    pub async fn keygen(&self) -> Result<AccountOwner> {
1287        let stdout = self
1288            .command()
1289            .await?
1290            .arg("keygen")
1291            .spawn_and_wait_for_stdout()
1292            .await?;
1293        AccountOwner::from_str(stdout.as_str().trim())
1294    }
1295
1296    /// Returns the default chain.
1297    pub fn default_chain(&self) -> Option<ChainId> {
1298        self.load_wallet().ok()?.default_chain()
1299    }
1300
1301    /// Runs `linera assign`.
1302    pub async fn assign(&self, owner: AccountOwner, chain_id: ChainId) -> Result<()> {
1303        let _stdout = self
1304            .command()
1305            .await?
1306            .arg("assign")
1307            .args(["--owner", &owner.to_string()])
1308            .args(["--chain-id", &chain_id.to_string()])
1309            .spawn_and_wait_for_stdout()
1310            .await?;
1311        Ok(())
1312    }
1313
1314    /// Runs `linera set-preferred-owner` for `chain_id`.
1315    pub async fn set_preferred_owner(
1316        &self,
1317        chain_id: ChainId,
1318        owner: Option<AccountOwner>,
1319    ) -> Result<()> {
1320        let mut owner_arg = vec!["--owner".to_string()];
1321        if let Some(owner) = owner {
1322            owner_arg.push(owner.to_string());
1323        };
1324        self.command()
1325            .await?
1326            .arg("set-preferred-owner")
1327            .args(["--chain-id", &chain_id.to_string()])
1328            .args(owner_arg)
1329            .spawn_and_wait_for_stdout()
1330            .await?;
1331        Ok(())
1332    }
1333
1334    /// Builds the application at `path` with `cargo` and returns its contract and service paths.
1335    pub async fn build_application(
1336        &self,
1337        path: &Path,
1338        name: &str,
1339        is_workspace: bool,
1340    ) -> Result<(PathBuf, PathBuf)> {
1341        Command::new("cargo")
1342            .current_dir(self.path_provider.path())
1343            .arg("build")
1344            .arg("--release")
1345            .args(["--target", "wasm32-unknown-unknown"])
1346            .arg("--manifest-path")
1347            .arg(path.join("Cargo.toml"))
1348            .spawn_and_wait_for_stdout()
1349            .await?;
1350
1351        let release_dir = match is_workspace {
1352            true => path.join("../target/wasm32-unknown-unknown/release"),
1353            false => path.join("target/wasm32-unknown-unknown/release"),
1354        };
1355
1356        let contract = release_dir.join(format!("{}_contract.wasm", name.replace('-', "_")));
1357        let service = release_dir.join(format!("{}_service.wasm", name.replace('-', "_")));
1358
1359        let contract_size = fs_err::tokio::metadata(&contract).await?.len();
1360        let service_size = fs_err::tokio::metadata(&service).await?.len();
1361        info!("Done building application {name}: contract_size={contract_size}, service_size={service_size}");
1362
1363        Ok((contract, service))
1364    }
1365}
1366
1367impl Drop for ClientWrapper {
1368    fn drop(&mut self) {
1369        use std::process::Command as SyncCommand;
1370
1371        if self.on_drop != OnClientDrop::CloseChains {
1372            return;
1373        }
1374
1375        let Ok(binary_path) = self.binary_path.lock() else {
1376            error!("Failed to close chains because a thread panicked with a lock to `binary_path`");
1377            return;
1378        };
1379
1380        let Some(binary_path) = binary_path.as_ref() else {
1381            warn!(
1382                "Assuming no chains need to be closed, because the command binary was never \
1383                resolved and therefore presumably never called"
1384            );
1385            return;
1386        };
1387
1388        let working_directory = self.path_provider.path();
1389        let mut wallet_show_command = SyncCommand::new(binary_path);
1390
1391        for argument in self.command_arguments() {
1392            wallet_show_command.arg(&*argument);
1393        }
1394
1395        let Ok(wallet_show_output) = wallet_show_command
1396            .current_dir(working_directory)
1397            .args(["wallet", "show", "--short", "--owned"])
1398            .output()
1399        else {
1400            warn!("Failed to execute `wallet show --short` to list chains to close");
1401            return;
1402        };
1403
1404        if !wallet_show_output.status.success() {
1405            warn!("Failed to list chains in the wallet to close them");
1406            return;
1407        }
1408
1409        let Ok(chain_list_string) = String::from_utf8(wallet_show_output.stdout) else {
1410            warn!(
1411                "Failed to close chains because `linera wallet show --short` \
1412                returned a non-UTF-8 output"
1413            );
1414            return;
1415        };
1416
1417        let chain_ids = chain_list_string
1418            .split('\n')
1419            .map(|line| line.trim())
1420            .filter(|line| !line.is_empty());
1421
1422        for chain_id in chain_ids {
1423            let mut close_chain_command = SyncCommand::new(binary_path);
1424
1425            for argument in self.command_arguments() {
1426                close_chain_command.arg(&*argument);
1427            }
1428
1429            close_chain_command.current_dir(working_directory);
1430
1431            match close_chain_command.args(["close-chain", chain_id]).status() {
1432                Ok(status) if status.success() => (),
1433                Ok(failure) => warn!("Failed to close chain {chain_id}: {failure}"),
1434                Err(error) => warn!("Failed to close chain {chain_id}: {error}"),
1435            }
1436        }
1437    }
1438}
1439
1440#[cfg(with_testing)]
1441impl ClientWrapper {
1442    /// Builds the example application with the given name and returns its contract and service paths.
1443    pub async fn build_example(&self, name: &str) -> Result<(PathBuf, PathBuf)> {
1444        self.build_application(Self::example_path(name)?.as_path(), name, true)
1445            .await
1446    }
1447
1448    /// Returns the path to the example application with the given name.
1449    pub fn example_path(name: &str) -> Result<PathBuf> {
1450        Ok(env::current_dir()?.join("../examples/").join(name))
1451    }
1452}
1453
1454fn truncate_query_output(input: &str) -> String {
1455    let max_len = 1000;
1456    if input.len() < max_len {
1457        input.to_string()
1458    } else {
1459        format!("{} ...", input.get(..max_len).unwrap())
1460    }
1461}
1462
1463fn truncate_query_output_serialize<T: Serialize>(query: T) -> String {
1464    let query = serde_json::to_string(&query).expect("Failed to serialize the failed query");
1465    let max_len = 200;
1466    if query.len() < max_len {
1467        query
1468    } else {
1469        format!("{} ...", query.get(..max_len).unwrap())
1470    }
1471}
1472
1473/// A running node service.
1474pub struct NodeService {
1475    port: u16,
1476    child: Child,
1477}
1478
1479impl NodeService {
1480    fn new(port: u16, child: Child) -> Self {
1481        Self { port, child }
1482    }
1483
1484    /// Terminates the node service by killing its child process.
1485    pub async fn terminate(mut self) -> Result<()> {
1486        self.child.kill().await.context("terminating node service")
1487    }
1488
1489    /// Returns the port the node service is listening on.
1490    pub fn port(&self) -> u16 {
1491        self.port
1492    }
1493
1494    /// Checks that the node service child process is still running.
1495    pub fn ensure_is_running(&mut self) -> Result<()> {
1496        self.child.ensure_is_running()
1497    }
1498
1499    /// Issues a `processInbox` GraphQL mutation and returns the hashes of the created blocks.
1500    pub async fn process_inbox(&self, chain_id: &ChainId) -> Result<Vec<CryptoHash>> {
1501        let query = format!("mutation {{ processInbox(chainId: \"{chain_id}\") }}");
1502        let mut data = self.query_node(query).await?;
1503        Ok(serde_json::from_value(data["processInbox"].take())?)
1504    }
1505
1506    /// Issues a `sync` GraphQL mutation for the given chain and returns the new block height.
1507    pub async fn sync(&self, chain_id: &ChainId) -> Result<u64> {
1508        let query = format!("mutation {{ sync(chainId: \"{chain_id}\") }}");
1509        let mut data = self.query_node(query).await?;
1510        Ok(serde_json::from_value(data["sync"].take())?)
1511    }
1512
1513    /// Issues a `transfer` GraphQL mutation and returns the resulting block hash.
1514    pub async fn transfer(
1515        &self,
1516        chain_id: ChainId,
1517        owner: AccountOwner,
1518        recipient: Account,
1519        amount: Amount,
1520    ) -> Result<CryptoHash> {
1521        let json_owner = owner.to_value();
1522        let json_recipient = recipient.to_value();
1523        let query = format!(
1524            "mutation {{ transfer(\
1525                 chainId: \"{chain_id}\", \
1526                 owner: {json_owner}, \
1527                 recipient: {json_recipient}, \
1528                 amount: \"{amount}\") \
1529             }}"
1530        );
1531        let data = self.query_node(query).await?;
1532        serde_json::from_value(data["transfer"].clone())
1533            .context("missing transfer field in response")
1534    }
1535
1536    /// Queries the balance of the given account via GraphQL.
1537    pub async fn balance(&self, account: &Account) -> Result<Amount> {
1538        let chain = account.chain_id;
1539        let owner = account.owner;
1540        if matches!(owner, AccountOwner::CHAIN) {
1541            let query = format!(
1542                "query {{ chain(chainId:\"{chain}\") {{
1543                    executionState {{ system {{ balance }} }}
1544                }} }}"
1545            );
1546            let response = self.query_node(query).await?;
1547            let balance = &response["chain"]["executionState"]["system"]["balance"]
1548                .as_str()
1549                .unwrap();
1550            return Ok(Amount::from_str(balance)?);
1551        }
1552        let query = format!(
1553            "query {{ chain(chainId:\"{chain}\") {{
1554                executionState {{ system {{ balances {{
1555                    entry(key:\"{owner}\") {{ value }}
1556                }} }} }}
1557            }} }}"
1558        );
1559        let response = self.query_node(query).await?;
1560        let balances = &response["chain"]["executionState"]["system"]["balances"];
1561        let balance = balances["entry"]["value"].as_str();
1562        match balance {
1563            None => Ok(Amount::ZERO),
1564            Some(amount) => Ok(Amount::from_str(amount)?),
1565        }
1566    }
1567
1568    /// Returns an [`ApplicationWrapper`] for querying the given application on this node service.
1569    pub fn make_application<A: ContractAbi>(
1570        &self,
1571        chain_id: &ChainId,
1572        application_id: &ApplicationId<A>,
1573    ) -> Result<ApplicationWrapper<A>> {
1574        let application_id = application_id.forget_abi().to_string();
1575        let link = format!(
1576            "http://localhost:{}/chains/{chain_id}/applications/{application_id}",
1577            self.port
1578        );
1579        Ok(ApplicationWrapper::from(link))
1580    }
1581
1582    /// Issues a `publishDataBlob` GraphQL mutation and returns the blob's hash.
1583    pub async fn publish_data_blob(
1584        &self,
1585        chain_id: &ChainId,
1586        bytes: Vec<u8>,
1587    ) -> Result<CryptoHash> {
1588        let query = format!(
1589            "mutation {{ publishDataBlob(chainId: {}, bytes: {}) }}",
1590            chain_id.to_value(),
1591            bytes.to_value(),
1592        );
1593        let data = self.query_node(query).await?;
1594        serde_json::from_value(data["publishDataBlob"].clone())
1595            .context("missing publishDataBlob field in response")
1596    }
1597
1598    /// Issues a `publishModule` GraphQL mutation and returns the new module ID.
1599    pub async fn publish_module<Abi, Parameters, InstantiationArgument>(
1600        &self,
1601        chain_id: &ChainId,
1602        contract: PathBuf,
1603        service: PathBuf,
1604        vm_runtime: VmRuntime,
1605    ) -> Result<ModuleId<Abi, Parameters, InstantiationArgument>> {
1606        let contract_code = Bytecode::load_from_file(&contract)?;
1607        let service_code = Bytecode::load_from_file(&service)?;
1608        let query = format!(
1609            "mutation {{ publishModule(chainId: {}, contract: {}, service: {}, vmRuntime: {}) }}",
1610            chain_id.to_value(),
1611            contract_code.to_value(),
1612            service_code.to_value(),
1613            vm_runtime.to_value(),
1614        );
1615        let data = self.query_node(query).await?;
1616        let module_str = data["publishModule"]
1617            .as_str()
1618            .context("module ID not found")?;
1619        let module_id: ModuleId = module_str.parse().context("could not parse module ID")?;
1620        Ok(module_id.with_abi())
1621    }
1622
1623    /// Queries the committees of the given chain via GraphQL.
1624    pub async fn query_committees(&self, chain_id: &ChainId) -> Result<BTreeMap<Epoch, Committee>> {
1625        let query = format!(
1626            "query {{ chain(chainId:\"{chain_id}\") {{
1627                executionState {{ system {{ committees }} }}
1628            }} }}"
1629        );
1630        let mut response = self.query_node(query).await?;
1631        let committees = response["chain"]["executionState"]["system"]["committees"].take();
1632        Ok(serde_json::from_value(committees)?)
1633    }
1634
1635    /// Issues an `eventsFromIndex` GraphQL query for events on the given chain and stream.
1636    pub async fn events_from_index(
1637        &self,
1638        chain_id: &ChainId,
1639        stream_id: &StreamId,
1640        start_index: u32,
1641    ) -> Result<Vec<IndexAndEvent>> {
1642        let query = format!(
1643            "query {{
1644               eventsFromIndex(chainId: \"{chain_id}\", streamId: {}, startIndex: {start_index})
1645               {{ index event }}
1646             }}",
1647            stream_id.to_value()
1648        );
1649        let mut response = self.query_node(query).await?;
1650        let response = response["eventsFromIndex"].take();
1651        Ok(serde_json::from_value(response)?)
1652    }
1653
1654    /// Posts the given GraphQL query to the node service, retrying on timeouts and errors.
1655    pub async fn query_node(&self, query: impl AsRef<str>) -> Result<Value> {
1656        let n_try = 5;
1657        let query = query.as_ref();
1658        for i in 0..n_try {
1659            linera_base::time::timer::sleep(Duration::from_secs(i)).await;
1660            let url = format!("http://localhost:{}/", self.port);
1661            let client = reqwest_client();
1662            let result = client
1663                .post(url)
1664                .json(&json!({ "query": query }))
1665                .send()
1666                .await;
1667            if matches!(result, Err(ref error) if error.is_timeout()) {
1668                warn!(
1669                    "Timeout when sending query {} to the node service",
1670                    truncate_query_output(query)
1671                );
1672                continue;
1673            }
1674            let response = result.with_context(|| {
1675                format!(
1676                    "query_node: failed to post query={}",
1677                    truncate_query_output(query)
1678                )
1679            })?;
1680            ensure!(
1681                response.status().is_success(),
1682                "Query \"{}\" failed: {}",
1683                truncate_query_output(query),
1684                response
1685                    .text()
1686                    .await
1687                    .unwrap_or_else(|error| format!("Could not get response text: {error}"))
1688            );
1689            let value: Value = response.json().await.context("invalid JSON")?;
1690            if let Some(errors) = value.get("errors") {
1691                warn!(
1692                    "Query \"{}\" failed: {}",
1693                    truncate_query_output(query),
1694                    errors
1695                );
1696            } else {
1697                return Ok(value["data"].clone());
1698            }
1699        }
1700        bail!(
1701            "Query \"{}\" failed after {} retries.",
1702            truncate_query_output(query),
1703            n_try
1704        );
1705    }
1706
1707    /// Issues a `createApplication` GraphQL mutation and returns the new application ID.
1708    pub async fn create_application<
1709        Abi: ContractAbi,
1710        Parameters: Serialize,
1711        InstantiationArgument: Serialize,
1712    >(
1713        &self,
1714        chain_id: &ChainId,
1715        module_id: &ModuleId<Abi, Parameters, InstantiationArgument>,
1716        parameters: &Parameters,
1717        argument: &InstantiationArgument,
1718        required_application_ids: &[ApplicationId],
1719    ) -> Result<ApplicationId<Abi>> {
1720        let module_id = module_id.forget_abi();
1721        let json_required_applications_ids = required_application_ids
1722            .iter()
1723            .map(ApplicationId::to_string)
1724            .collect::<Vec<_>>()
1725            .to_value();
1726        // Convert to `serde_json::Value` then `async_graphql::Value` via the trait `InputType`.
1727        let new_parameters = serde_json::to_value(parameters)
1728            .context("could not create parameters JSON")?
1729            .to_value();
1730        let new_argument = serde_json::to_value(argument)
1731            .context("could not create argument JSON")?
1732            .to_value();
1733        let query = format!(
1734            "mutation {{ createApplication(\
1735                 chainId: \"{chain_id}\",
1736                 moduleId: \"{module_id}\", \
1737                 parameters: {new_parameters}, \
1738                 instantiationArgument: {new_argument}, \
1739                 requiredApplicationIds: {json_required_applications_ids}) \
1740             }}"
1741        );
1742        let data = self.query_node(query).await?;
1743        let app_id_str = data["createApplication"]
1744            .as_str()
1745            .context("missing createApplication string in response")?
1746            .trim();
1747        Ok(app_id_str
1748            .parse::<ApplicationId>()
1749            .context("invalid application ID")?
1750            .with_abi())
1751    }
1752
1753    /// Obtains the hash and height of the `chain`'s tip block, as known by this node service.
1754    pub async fn chain_tip(&self, chain: ChainId) -> Result<Option<(CryptoHash, BlockHeight)>> {
1755        let query = format!(
1756            r#"query {{ block(chainId: "{chain}") {{
1757                hash
1758                block {{ header {{ height }} }}
1759            }} }}"#
1760        );
1761
1762        let mut response = self.query_node(&query).await?;
1763
1764        match (
1765            mem::take(&mut response["block"]["hash"]),
1766            mem::take(&mut response["block"]["block"]["header"]["height"]),
1767        ) {
1768            (Value::Null, Value::Null) => Ok(None),
1769            (Value::String(hash), Value::Number(height)) => Ok(Some((
1770                hash.parse()
1771                    .context("Received an invalid hash {hash:?} for chain tip")?,
1772                BlockHeight(height.as_u64().unwrap()),
1773            ))),
1774            invalid_data => bail!("Expected a tip hash string, but got {invalid_data:?} instead"),
1775        }
1776    }
1777
1778    /// Subscribes to the node service and returns a stream of notifications about a chain.
1779    pub async fn notifications(
1780        &self,
1781        chain_id: ChainId,
1782    ) -> Result<Pin<Box<impl Stream<Item = Result<Notification>>>>> {
1783        let query = format!("subscription {{ notifications(chainId: \"{chain_id}\") }}",);
1784        let url = format!("ws://localhost:{}/ws", self.port);
1785        let mut request = url.into_client_request()?;
1786        request.headers_mut().insert(
1787            "Sec-WebSocket-Protocol",
1788            HeaderValue::from_str("graphql-transport-ws")?,
1789        );
1790        let (mut websocket, _) = async_tungstenite::tokio::connect_async(request).await?;
1791        let init_json = json!({
1792          "type": "connection_init",
1793          "payload": {}
1794        });
1795        websocket.send(init_json.to_string().into()).await?;
1796        let text = websocket
1797            .next()
1798            .await
1799            .context("Failed to establish connection")??
1800            .into_text()?;
1801        ensure!(
1802            text == "{\"type\":\"connection_ack\"}",
1803            "Unexpected response: {text}"
1804        );
1805        let query_json = json!({
1806          "id": "1",
1807          "type": "start",
1808          "payload": {
1809            "query": query,
1810            "variables": {},
1811            "operationName": null
1812          }
1813        });
1814        websocket.send(query_json.to_string().into()).await?;
1815        Ok(Box::pin(websocket.map_err(anyhow::Error::from).and_then(
1816            |message| async {
1817                let text = message.into_text()?;
1818                let value: Value = serde_json::from_str(&text).context("invalid JSON")?;
1819                if let Some(errors) = value["payload"].get("errors") {
1820                    bail!("Notification subscription failed: {errors:?}");
1821                }
1822                serde_json::from_value(value["payload"]["data"]["notifications"].clone())
1823                    .context("Failed to deserialize notification")
1824            },
1825        )))
1826    }
1827
1828    /// Subscribes to query results via the `queryResult` GraphQL subscription.
1829    pub async fn query_result(
1830        &self,
1831        name: &str,
1832        chain_id: ChainId,
1833        application_id: &ApplicationId,
1834    ) -> Result<Pin<Box<impl Stream<Item = Result<Value>>>>> {
1835        let query = format!(
1836            r#"subscription {{ queryResult(name: "{name}", chainId: "{chain_id}", applicationId: "{application_id}") }}"#,
1837        );
1838        let url = format!("ws://localhost:{}/ws", self.port);
1839        let mut request = url.into_client_request()?;
1840        request.headers_mut().insert(
1841            "Sec-WebSocket-Protocol",
1842            HeaderValue::from_str("graphql-transport-ws")?,
1843        );
1844        let (mut websocket, _) = async_tungstenite::tokio::connect_async(request).await?;
1845        let init_json = json!({
1846          "type": "connection_init",
1847          "payload": {}
1848        });
1849        websocket.send(init_json.to_string().into()).await?;
1850        let text = websocket
1851            .next()
1852            .await
1853            .context("Failed to establish connection")??
1854            .into_text()?;
1855        ensure!(
1856            text == "{\"type\":\"connection_ack\"}",
1857            "Unexpected response: {text}"
1858        );
1859        let query_json = json!({
1860          "id": "1",
1861          "type": "start",
1862          "payload": {
1863            "query": query,
1864            "variables": {},
1865            "operationName": null
1866          }
1867        });
1868        websocket.send(query_json.to_string().into()).await?;
1869        Ok(Box::pin(websocket.map_err(anyhow::Error::from).and_then(
1870            |message| async {
1871                let text = message.into_text()?;
1872                let value: Value = serde_json::from_str(&text).context("invalid JSON")?;
1873                if let Some(errors) = value["payload"].get("errors") {
1874                    bail!("Query result subscription failed: {errors:?}");
1875                }
1876                Ok(value["payload"]["data"]["queryResult"].clone())
1877            },
1878        )))
1879    }
1880}
1881
1882/// A running faucet service.
1883pub struct FaucetService {
1884    port: u16,
1885    child: Child,
1886    _temp_dir: tempfile::TempDir,
1887}
1888
1889impl FaucetService {
1890    fn new(port: u16, child: Child, temp_dir: tempfile::TempDir) -> Self {
1891        Self {
1892            port,
1893            child,
1894            _temp_dir: temp_dir,
1895        }
1896    }
1897
1898    /// Terminates the faucet service by killing its child process.
1899    pub async fn terminate(mut self) -> Result<()> {
1900        self.child
1901            .kill()
1902            .await
1903            .context("terminating faucet service")
1904    }
1905
1906    /// Checks that the faucet service child process is still running.
1907    pub fn ensure_is_running(&mut self) -> Result<()> {
1908        self.child.ensure_is_running()
1909    }
1910
1911    /// Returns a [`Faucet`] client connected to this running faucet service.
1912    pub fn instance(&self) -> Faucet {
1913        Faucet::new(format!("http://localhost:{}/", self.port))
1914    }
1915}
1916
1917/// A running `Application` to be queried in GraphQL.
1918pub struct ApplicationWrapper<A> {
1919    uri: String,
1920    _phantom: PhantomData<A>,
1921}
1922
1923impl<A> ApplicationWrapper<A> {
1924    /// Runs the given GraphQL query against the application and returns the response data.
1925    pub async fn run_graphql_query(&self, query: impl AsRef<str>) -> Result<Value> {
1926        let query = query.as_ref();
1927        let value = self.run_json_query(json!({ "query": query })).await?;
1928        Ok(value["data"].clone())
1929    }
1930
1931    /// Posts the given serializable JSON query to the application endpoint, retrying on failure.
1932    pub async fn run_json_query<T: Serialize>(&self, query: T) -> Result<Value> {
1933        const MAX_RETRIES: usize = 5;
1934
1935        for i in 0.. {
1936            let client = reqwest_client();
1937            let result = client.post(&self.uri).json(&query).send().await;
1938            let response = match result {
1939                Ok(response) => response,
1940                Err(error) if i < MAX_RETRIES => {
1941                    warn!(
1942                        "Failed to post query \"{}\": {error}; retrying",
1943                        truncate_query_output_serialize(&query),
1944                    );
1945                    continue;
1946                }
1947                Err(error) => {
1948                    let query = truncate_query_output_serialize(&query);
1949                    return Err(error)
1950                        .with_context(|| format!("run_json_query: failed to post query={query}"));
1951                }
1952            };
1953            ensure!(
1954                response.status().is_success(),
1955                "Query \"{}\" failed: {}",
1956                truncate_query_output_serialize(&query),
1957                response
1958                    .text()
1959                    .await
1960                    .unwrap_or_else(|error| format!("Could not get response text: {error}"))
1961            );
1962            let value: Value = response.json().await.context("invalid JSON")?;
1963            if let Some(errors) = value.get("errors") {
1964                bail!(
1965                    "Query \"{}\" failed: {}",
1966                    truncate_query_output_serialize(&query),
1967                    errors
1968                );
1969            }
1970            return Ok(value);
1971        }
1972        unreachable!()
1973    }
1974
1975    /// Runs the given string as a GraphQL `query` and returns the response data.
1976    pub async fn query(&self, query: impl AsRef<str>) -> Result<Value> {
1977        let query = query.as_ref();
1978        self.run_graphql_query(&format!("query {{ {query} }}"))
1979            .await
1980    }
1981
1982    /// Runs the given string as a GraphQL `query` and deserializes the named field's value.
1983    pub async fn query_json<T: DeserializeOwned>(&self, query: impl AsRef<str>) -> Result<T> {
1984        let query = query.as_ref().trim();
1985        let name = query
1986            .split_once(|ch: char| !ch.is_alphanumeric())
1987            .map_or(query, |(name, _)| name);
1988        let data = self.query(query).await?;
1989        serde_json::from_value(data[name].clone())
1990            .with_context(|| format!("{name} field missing in response"))
1991    }
1992
1993    /// Runs the given string as a GraphQL `mutation` and returns the response data.
1994    pub async fn mutate(&self, mutation: impl AsRef<str>) -> Result<Value> {
1995        let mutation = mutation.as_ref();
1996        self.run_graphql_query(&format!("mutation {{ {mutation} }}"))
1997            .await
1998    }
1999
2000    /// Runs several GraphQL mutations in a single aliased `mutation` request.
2001    pub async fn multiple_mutate(&self, mutations: &[String]) -> Result<Value> {
2002        let mut out = String::from("mutation {\n");
2003        for (index, mutation) in mutations.iter().enumerate() {
2004            out = format!("{out}  u{index}: {mutation}\n");
2005        }
2006        out.push_str("}\n");
2007        self.run_graphql_query(&out).await
2008    }
2009}
2010
2011impl<A> From<String> for ApplicationWrapper<A> {
2012    fn from(uri: String) -> ApplicationWrapper<A> {
2013        ApplicationWrapper {
2014            uri,
2015            _phantom: PhantomData,
2016        }
2017    }
2018}
2019
2020/// Returns the timeout for tests that wait for notifications, either read from the env
2021/// variable `LINERA_TEST_NOTIFICATION_TIMEOUT_MS`, or the default value of 10 seconds.
2022#[cfg(with_testing)]
2023fn notification_timeout() -> Duration {
2024    const NOTIFICATION_TIMEOUT_MS_ENV: &str = "LINERA_TEST_NOTIFICATION_TIMEOUT_MS";
2025    const NOTIFICATION_TIMEOUT_MS_DEFAULT: u64 = 10_000;
2026
2027    match env::var(NOTIFICATION_TIMEOUT_MS_ENV) {
2028        Ok(var) => Duration::from_millis(var.parse().unwrap_or_else(|error| {
2029            panic!("{NOTIFICATION_TIMEOUT_MS_ENV} is not a valid number: {error}")
2030        })),
2031        Err(env::VarError::NotPresent) => Duration::from_millis(NOTIFICATION_TIMEOUT_MS_DEFAULT),
2032        Err(env::VarError::NotUnicode(_)) => {
2033            panic!("{NOTIFICATION_TIMEOUT_MS_ENV} must be valid Unicode")
2034        }
2035    }
2036}
2037
2038#[cfg(with_testing)]
2039/// Extension trait for streams of [`Notification`]s, providing helpers to wait for events.
2040pub trait NotificationsExt {
2041    /// Waits for a notification for which `f` returns `Some(t)`, and returns `t`.
2042    fn wait_for<T>(
2043        &mut self,
2044        f: impl FnMut(Notification) -> Option<T>,
2045    ) -> impl Future<Output = Result<T>>;
2046
2047    /// Waits for a `NewEvents` notification for the given block height. If no height is specified,
2048    /// any height is accepted.
2049    fn wait_for_events(
2050        &mut self,
2051        expected_height: impl Into<Option<BlockHeight>>,
2052    ) -> impl Future<Output = Result<BTreeSet<StreamId>>> {
2053        let expected_height = expected_height.into();
2054        self.wait_for(move |notification| {
2055            if let Reason::NewEvents {
2056                height,
2057                event_streams,
2058                ..
2059            } = notification.reason
2060            {
2061                if expected_height.is_none_or(|h| h == height) && !event_streams.is_empty() {
2062                    return Some(event_streams);
2063                }
2064            }
2065            None
2066        })
2067    }
2068
2069    /// Waits for a `NewBlock` notification for the given block height. If no height is specified,
2070    /// any height is accepted.
2071    fn wait_for_block(
2072        &mut self,
2073        expected_height: impl Into<Option<BlockHeight>>,
2074    ) -> impl Future<Output = Result<CryptoHash>> {
2075        let expected_height = expected_height.into();
2076        self.wait_for(move |notification| {
2077            if let Reason::NewBlock { height, hash, .. } = notification.reason {
2078                if expected_height.is_none_or(|h| h == height) {
2079                    return Some(hash);
2080                }
2081            }
2082            None
2083        })
2084    }
2085
2086    /// Waits for a `NewIncomingBundle` notification for the given sender chain and sender block
2087    /// height. If no height is specified, any height is accepted.
2088    fn wait_for_bundle(
2089        &mut self,
2090        expected_origin: ChainId,
2091        expected_height: impl Into<Option<BlockHeight>>,
2092    ) -> impl Future<Output = Result<()>> {
2093        let expected_height = expected_height.into();
2094        self.wait_for(move |notification| {
2095            if let Reason::NewIncomingBundle { height, origin } = notification.reason {
2096                if expected_height.is_none_or(|h| h == height) && origin == expected_origin {
2097                    return Some(());
2098                }
2099            }
2100            None
2101        })
2102    }
2103}
2104
2105#[cfg(with_testing)]
2106impl<S: Stream<Item = Result<Notification>>> NotificationsExt for Pin<Box<S>> {
2107    async fn wait_for<T>(&mut self, mut f: impl FnMut(Notification) -> Option<T>) -> Result<T> {
2108        let mut timeout = Box::pin(linera_base::time::timer::sleep(notification_timeout())).fuse();
2109        loop {
2110            let notification = futures::select! {
2111                () = timeout => bail!("Timeout waiting for notification"),
2112                notification = self.next().fuse() => notification.context("Stream closed")??,
2113            };
2114            if let Some(t) = f(notification) {
2115                return Ok(t);
2116            }
2117        }
2118    }
2119}