Skip to main content

iota_single_node_benchmark/
mock_account.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use std::{collections::BTreeMap, sync::Arc};
6
7use futures::stream::FuturesUnordered;
8use iota_sdk_types::{Address, ObjectReference};
9use iota_types::{
10    crypto::{AccountPrivateKey, get_account_private_key},
11    object::Object,
12};
13
14#[derive(Clone)]
15pub struct Account {
16    pub sender: Address,
17    pub private_key: Arc<AccountPrivateKey>,
18    pub gas_objects: Arc<Vec<ObjectReference>>,
19}
20
21/// Generate \num_accounts accounts and for each account generate
22/// \gas_object_num_per_account gas objects. Return all accounts along with a
23/// flattened list of all gas objects as genesis objects.
24pub async fn batch_create_account_and_gas(
25    num_accounts: u64,
26    gas_object_num_per_account: u64,
27) -> (BTreeMap<Address, Account>, Vec<Object>) {
28    let tasks: FuturesUnordered<_> = (0..num_accounts)
29        .map(|_| {
30            tokio::spawn(async move {
31                let (sender, private_key) = get_account_private_key();
32                let objects = (0..gas_object_num_per_account)
33                    .map(|_| Object::with_owner_for_testing(sender))
34                    .collect::<Vec<_>>();
35                (sender, private_key, objects)
36            })
37        })
38        .collect();
39    let mut accounts = BTreeMap::new();
40    let mut genesis_gas_objects = vec![];
41    for task in tasks {
42        let (sender, private_key, gas_objects) = task.await.unwrap();
43        let gas_object_refs: Vec<_> = gas_objects.iter().map(|o| o.object_ref()).collect();
44        accounts.insert(
45            sender,
46            Account {
47                sender,
48                private_key: Arc::new(private_key),
49                gas_objects: Arc::new(gas_object_refs),
50            },
51        );
52        genesis_gas_objects.extend(gas_objects);
53    }
54    (accounts, genesis_gas_objects)
55}