Skip to content

Commit

Permalink
fix(test): use independent client instances to reduce load
Browse files Browse the repository at this point in the history
  • Loading branch information
RolandSherwin committed Jul 5, 2024
1 parent a1c9f0d commit 85d53d6
Show file tree
Hide file tree
Showing 2 changed files with 27 additions and 7 deletions.
7 changes: 7 additions & 0 deletions sn_node/tests/common/client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
27 changes: 20 additions & 7 deletions sn_node/tests/spend_simulation.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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(),
);
Expand Down Expand Up @@ -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(())
}
Expand Down Expand Up @@ -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"));
Expand Down Expand Up @@ -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(),
Expand Down Expand Up @@ -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!(
Expand All @@ -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.
Expand Down

0 comments on commit 85d53d6

Please sign in to comment.