Add unit test for saving wallet with pending blobs

linera-client/Cargo.toml

@@ -120,6 +120,7 @@ non-fungible.workspace = true
 proptest.workspace = true
 reqwest = { workspace = true, features = ["json"] }
 social.workspace = true
+tempfile.workspace = true
 test-case.workspace = true
 test-log = { workspace = true, features = ["trace"] }
 test-strategy.workspace = true

linera-client/src/chain_listener.rs

@@ -26,10 +26,6 @@ use tracing::{debug, error, info, warn, Instrument as _};
 
 use crate::{wallet::Wallet, Error};
 
-#[cfg(test)]
-#[path = "unit_tests/chain_listener.rs"]
-mod tests;
-
 #[derive(Debug, Default, Clone, clap::Args)]
 pub struct ChainListenerConfig {
     /// Do not create blocks automatically to receive incoming messages. Instead, wait for

linera-client/src/client_context.rs

@@ -10,8 +10,8 @@ use async_trait::async_trait;
 use futures::Future;
 use linera_base::{
     crypto::KeyPair,
-    data_types::{BlockHeight, Timestamp},
-    identifiers::{Account, ChainId},
+    data_types::{Blob, BlockHeight, Timestamp},
+    identifiers::{Account, BlobId, ChainId},
     ownership::ChainOwnership,
     time::{Duration, Instant},
 };
@@ -184,6 +184,41 @@ where
         }
     }
 
+    #[cfg(with_testing)]
+    pub fn new_test(storage: S, wallet: W) -> Self {
+        let send_recv_timeout = Duration::from_millis(4000);
+        let retry_delay = Duration::from_millis(1000);
+        let max_retries = 10;
+
+        let node_options = NodeOptions {
+            send_timeout: send_recv_timeout,
+            recv_timeout: send_recv_timeout,
+            retry_delay,
+            max_retries,
+        };
+        let node_provider = NodeProvider::new(node_options);
+        let delivery = CrossChainMessageDelivery::new(true);
+        let chain_ids = wallet.chain_ids();
+        let name = match chain_ids.len() {
+            0 => "Client node".to_string(),
+            1 => format!("Client node for {:.8}", chain_ids[0]),
+            n => format!("Client node for {:.8} and {} others", chain_ids[0], n - 1),
+        };
+        let client = Client::new(node_provider, storage, 10, delivery, false, chain_ids, name);
+
+        ClientContext {
+            client: Arc::new(client),
+            wallet: WalletState::new(wallet),
+            send_timeout: send_recv_timeout,
+            recv_timeout: send_recv_timeout,
+            retry_delay,
+            max_retries,
+            chain_listeners: JoinSet::default(),
+            blanket_message_policy: BlanketMessagePolicy::Accept,
+            restrict_chain_ids_to: None,
+        }
+    }
+
     /// Retrieve the default account. Current this is the common account of the default
     /// chain.
     pub fn default_account(&self) -> Account {
@@ -268,6 +303,29 @@ where
         chain_id: ChainId,
         key_pair: Option<KeyPair>,
         timestamp: Timestamp,
+    ) -> Result<(), Error> {
+        self.update_wallet_for_new_chain_internal(chain_id, key_pair, timestamp, BTreeMap::new())
+            .await
+    }
+
+    #[cfg(test)]
+    pub async fn update_wallet_for_new_chain_with_pending_blobs(
+        &mut self,
+        chain_id: ChainId,
+        key_pair: Option<KeyPair>,
+        timestamp: Timestamp,
+        pending_blobs: BTreeMap<BlobId, Blob>,
+    ) -> Result<(), Error> {
+        self.update_wallet_for_new_chain_internal(chain_id, key_pair, timestamp, pending_blobs)
+            .await
+    }
+
+    async fn update_wallet_for_new_chain_internal(
+        &mut self,
+        chain_id: ChainId,
+        key_pair: Option<KeyPair>,
+        timestamp: Timestamp,
+        pending_blobs: BTreeMap<BlobId, Blob>,
     ) -> Result<(), Error> {
         if self.wallet.get(chain_id).is_none() {
             self.mutate_wallet(|w| {
@@ -278,7 +336,7 @@ where
                     timestamp,
                     next_block_height: BlockHeight::ZERO,
                     pending_block: None,
-                    pending_blobs: BTreeMap::new(),
+                    pending_blobs,
                 })
             })
             .await?;

linera-client/src/lib.rs

@@ -16,4 +16,7 @@ pub mod storage;
 pub mod util;
 pub mod wallet;
 
+#[cfg(test)]
+mod unit_tests;
+
 pub use error::Error;

linera-client/src/unit_tests/chain_listener.rs

@@ -5,6 +5,7 @@
 
 use std::{collections::BTreeMap, sync::Arc};
 
+use super::util::make_genesis_config;
 use async_trait::async_trait;
 use futures::{lock::Mutex, FutureExt as _};
 use linera_base::{
@@ -18,18 +19,13 @@ use linera_core::{
     node::CrossChainMessageDelivery,
     test_utils::{MemoryStorageBuilder, NodeProvider, StorageBuilder as _, TestBuilder},
 };
-use linera_execution::{system::Recipient, ResourceControlPolicy};
-use linera_rpc::{
-    config::{NetworkProtocol, ValidatorPublicNetworkPreConfig},
-    simple::TransportProtocol,
-};
+use linera_execution::system::Recipient;
 use linera_storage::{DbStorage, TestClock};
 use linera_views::memory::MemoryStore;
 use rand::SeedableRng as _;
 
 use crate::{
     chain_listener::{self, ChainListener, ChainListenerConfig, ClientContext as _},
-    config::{CommitteeConfig, GenesisConfig, ValidatorConfig},
     wallet::{UserChain, Wallet},
     Error,
 };
@@ -108,30 +104,6 @@ impl chain_listener::ClientContext for ClientContext {
     }
 }
 
-fn make_genesis_config(builder: &TestBuilder<MemoryStorageBuilder>) -> GenesisConfig {
-    let network = ValidatorPublicNetworkPreConfig {
-        protocol: NetworkProtocol::Simple(TransportProtocol::Tcp),
-        host: "localhost".to_string(),
-        port: 8080,
-    };
-    let validator_names = builder.initial_committee.validators().keys();
-    let validators = validator_names
-        .map(|name| ValidatorConfig {
-            name: *name,
-            network: network.clone(),
-        })
-        .collect();
-    let mut genesis_config = GenesisConfig::new(
-        CommitteeConfig { validators },
-        builder.admin_id(),
-        Timestamp::from(0),
-        ResourceControlPolicy::default(),
-        "test network".to_string(),
-    );
-    genesis_config.chains.extend(builder.genesis_chains());
-    genesis_config
-}
-
 /// Tests that the chain listener, if there is a message in the inbox, will continue requesting
 /// timeout certificates until it becomes the leader and can process the inbox.
 #[test_log::test(tokio::test)]

linera-client/src/unit_tests/mod.rs

@@ -0,0 +1,11 @@
+// Copyright (c) Zefchain Labs, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+#[cfg(test)]
+mod util;
+
+#[cfg(test)]
+mod chain_listener;
+
+#[cfg(test)]
+mod wallet;

linera-client/src/unit_tests/util.rs

@@ -0,0 +1,36 @@
+// Copyright (c) Zefchain Labs, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+use linera_base::data_types::Timestamp;
+use linera_core::test_utils::{MemoryStorageBuilder, TestBuilder};
+use linera_execution::ResourceControlPolicy;
+use linera_rpc::{
+    config::{NetworkProtocol, ValidatorPublicNetworkPreConfig},
+    simple::TransportProtocol,
+};
+
+use crate::config::{CommitteeConfig, GenesisConfig, ValidatorConfig};
+
+pub fn make_genesis_config(builder: &TestBuilder<MemoryStorageBuilder>) -> GenesisConfig {
+    let network = ValidatorPublicNetworkPreConfig {
+        protocol: NetworkProtocol::Simple(TransportProtocol::Tcp),
+        host: "localhost".to_string(),
+        port: 8080,
+    };
+    let validator_names = builder.initial_committee.validators().keys();
+    let validators = validator_names
+        .map(|name| ValidatorConfig {
+            name: *name,
+            network: network.clone(),
+        })
+        .collect();
+    let mut genesis_config = GenesisConfig::new(
+        CommitteeConfig { validators },
+        builder.admin_id(),
+        Timestamp::from(0),
+        ResourceControlPolicy::default(),
+        "test network".to_string(),
+    );
+    genesis_config.chains.extend(builder.genesis_chains());
+    genesis_config
+}

linera-client/src/unit_tests/wallet.rs

@@ -0,0 +1,61 @@
+// Copyright (c) Zefchain Labs, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+use std::collections::BTreeMap;
+
+use crate::{client_context::ClientContext, config::WalletState, wallet::Wallet};
+
+use super::util::make_genesis_config;
+use anyhow::anyhow;
+use linera_base::{
+    crypto::KeyPair,
+    data_types::{Amount, Blob},
+    identifiers::{ChainDescription, ChainId},
+};
+use linera_core::test_utils::{MemoryStorageBuilder, StorageBuilder, TestBuilder};
+use rand::SeedableRng as _;
+
+/// Tests wether we can correctly save a wallet that contains pending blobs.
+#[test_log::test(tokio::test)]
+async fn test_save_wallet_with_pending_blobs() -> anyhow::Result<()> {
+    let mut rng = rand::rngs::StdRng::seed_from_u64(42);
+    let storage_builder = MemoryStorageBuilder::default();
+    let clock = storage_builder.clock().clone();
+    let mut builder = TestBuilder::new(storage_builder, 4, 1).await?;
+    let description = ChainDescription::Root(0);
+    let chain_id = ChainId::from(description);
+    builder.add_initial_chain(description, Amount::ONE).await?;
+    let storage = builder.make_storage().await?;
+    let key_pair = KeyPair::generate_from(&mut rng);
+
+    let genesis_config = make_genesis_config(&builder);
+
+    let tmp_dir = tempfile::tempdir()?;
+    let mut config_dir = tmp_dir.into_path();
+    config_dir.push("linera");
+    if !config_dir.exists() {
+        tracing::debug!("{} does not exist, creating", config_dir.display());
+        fs_err::create_dir(&config_dir)?;
+        tracing::debug!("{} created.", config_dir.display());
+    }
+    let wallet_path = config_dir.join("wallet.json");
+    if wallet_path.exists() {
+        return Err(anyhow!("Wallet already exists!"));
+    }
+    let wallet =
+        WalletState::create_from_file(&wallet_path, Wallet::new(genesis_config, Some(37)))?;
+    let mut context = ClientContext::new_test(storage, wallet);
+    let blob = Blob::new_data(b"blob".to_vec());
+    let mut pending_blobs = BTreeMap::new();
+    pending_blobs.insert(blob.id(), blob);
+    context
+        .update_wallet_for_new_chain_with_pending_blobs(
+            chain_id,
+            Some(key_pair),
+            clock.current_time(),
+            pending_blobs,
+        )
+        .await?;
+    context.save_wallet().await?;
+    Ok(())
+}