From 85d53d68f47c5f236087da736e0893a1c34f2cea Mon Sep 17 00:00:00 2001
From: Roland Sherwin <RolandSherwin@protonmail.com>
Date: Sat, 6 Jul 2024 02:53:50 +0530
Subject: [PATCH] fix(test): use independent client instances to reduce load

---
 sn_node/tests/common/client.rs    |  7 +++++++
 sn_node/tests/spend_simulation.rs | 27 ++++++++++++++++++++-------
 2 files changed, 27 insertions(+), 7 deletions(-)

diff --git a/sn_node/tests/common/client.rs b/sn_node/tests/common/client.rs
index e00ff36058..f7bc44b109 100644
--- a/sn_node/tests/common/client.rs
+++ b/sn_node/tests/common/client.rs
@@ -151,6 +151,13 @@ pub async fn get_client_and_funded_wallet(root_dir: &Path) -> Result<(Client, Ho
     }
 }
 
+pub async fn get_client() -> Client {
+    match DeploymentInventory::load() {
+        Ok(inventory) => Droplet::get_client(&inventory).await,
+        Err(_) => NonDroplet::get_client().await,
+    }
+}
+
 pub struct NonDroplet;
 impl NonDroplet {
     ///  Get a new Client for testing
diff --git a/sn_node/tests/spend_simulation.rs b/sn_node/tests/spend_simulation.rs
index f87371f684..154f0d10a1 100644
--- a/sn_node/tests/spend_simulation.rs
+++ b/sn_node/tests/spend_simulation.rs
@@ -9,7 +9,7 @@
 mod common;
 use assert_fs::TempDir;
 use assert_matches::assert_matches;
-use common::client::{get_client_and_funded_wallet, get_wallet};
+use common::client::{get_client, get_client_and_funded_wallet, get_wallet};
 use eyre::{bail, OptionExt, Report, Result};
 use itertools::Itertools;
 use rand::{seq::IteratorRandom, Rng};
@@ -94,6 +94,9 @@ struct State {
     action_senders: BTreeMap<WalletId, mpsc::Sender<WalletAction>>,
     /// The TempDir for each wallet. This has to be held until the end of the test.
     all_wallets: BTreeMap<WalletId, TempDir>,
+    /// Have a client per wallet. A single client instance can get overwhelmed with the number of transactions.
+    #[debug(skip)]
+    clients: BTreeMap<WalletId, Client>,
     /// The main pubkeys of all the wallets.
     main_pubkeys: BTreeMap<WalletId, MainPubkey>,
     /// The map from MainPubKey to WalletId. This is used to get wallets when we only have the cashnote in hand.
@@ -127,18 +130,23 @@ struct PendingTasksTracker {
 async fn spend_simulation() -> Result<()> {
     let _log_guards = LogBuilder::init_single_threaded_tokio_test("spend_simulation", true);
 
-    let (client, mut state) = init_state(MAX_WALLETS).await?;
+    let mut state = init_state(MAX_WALLETS).await?;
 
     let mut rng = rng::thread_rng();
     let (result_sender, mut result_rx) = mpsc::channel(10000);
 
     for (id, wallet_dir) in state.all_wallets.iter() {
+        let client = state
+            .clients
+            .get(id)
+            .ok_or_eyre("Client not found")?
+            .clone();
         let (action_sender, action_rx) = mpsc::channel(50);
         state.action_senders.insert(*id, action_sender);
         handle_action_per_wallet(
             *id,
             wallet_dir.to_path_buf(),
-            client.clone(),
+            client,
             action_rx,
             result_sender.clone(),
         );
@@ -237,7 +245,7 @@ async fn spend_simulation() -> Result<()> {
     info!("Final state: {state:?}. Sleeping before verifying wallets.");
     println!("Verifying all wallets in 10 seconds.");
     tokio::time::sleep(Duration::from_secs(10)).await;
-    verify_wallets(&state, client).await?;
+    verify_wallets(&state).await?;
 
     Ok(())
 }
@@ -493,8 +501,9 @@ async fn handle_wallet_task_result(
     Ok(())
 }
 
-async fn verify_wallets(state: &State, client: Client) -> Result<()> {
+async fn verify_wallets(state: &State) -> Result<()> {
     for (id, spends) in state.cashnotes_per_wallet.iter() {
+        let client = state.clients.get(id).ok_or_eyre("Client not found")?;
         println!("Verifying wallet {id}");
         info!("TestWallet {id} verifying {} spends", spends.len());
         let mut wallet = get_wallet(state.all_wallets.get(id).expect("Wallet not found"));
@@ -538,11 +547,12 @@ async fn verify_wallets(state: &State, client: Client) -> Result<()> {
 
 /// Create `count` number of wallets and fund them all with equal amounts of tokens.
 /// Return the client and the states of the wallets.
-async fn init_state(count: usize) -> Result<(Client, State)> {
+async fn init_state(count: usize) -> Result<State> {
     let mut state = State {
         all_wallets: BTreeMap::new(),
         main_pubkeys: BTreeMap::new(),
         action_senders: BTreeMap::new(),
+        clients: BTreeMap::new(),
         main_pubkeys_inverse: BTreeMap::new(),
         cashnote_tracker: BTreeMap::new(),
         cashnotes_per_wallet: BTreeMap::new(),
@@ -594,6 +604,9 @@ async fn init_state(count: usize) -> Result<(Client, State)> {
         .await?;
 
     for (id, address) in state.main_pubkeys.iter() {
+        // This sk is different from the wallet sk.
+        let client = get_client().await;
+        state.clients.insert(*id, client);
         let mut wallet = get_wallet(state.all_wallets.get(id).expect("Id should be present"));
         wallet.deposit_and_store_to_disk(&transfer.cash_notes_for_recipient)?;
         trace!(
@@ -618,7 +631,7 @@ async fn init_state(count: usize) -> Result<(Client, State)> {
         }
     }
 
-    Ok((client, state))
+    Ok(state)
 }
 
 /// Returns random recipients to send tokens to.