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_types::{
9    base_types::{IotaAddress, ObjectRef},
10    crypto::{AccountKeyPair, get_account_key_pair},
11    object::Object,
12};
13
14#[derive(Clone)]
15pub struct Account {
16    pub sender: IotaAddress,
17    pub keypair: Arc<AccountKeyPair>,
18    pub gas_objects: Arc<Vec<ObjectRef>>,
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<IotaAddress, Account>, Vec<Object>) {
28    let tasks: FuturesUnordered<_> = (0..num_accounts)
29        .map(|_| {
30            tokio::spawn(async move {
31                let (sender, keypair) = get_account_key_pair();
32                let objects = (0..gas_object_num_per_account)
33                    .map(|_| Object::with_owner_for_testing(sender))
34                    .collect::<Vec<_>>();
35                (sender, keypair, 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, keypair, gas_objects) = task.await.unwrap();
43        let gas_object_refs: Vec<_> = gas_objects
44            .iter()
45            .map(|o| o.compute_object_reference())
46            .collect();
47        accounts.insert(
48            sender,
49            Account {
50                sender,
51                keypair: Arc::new(keypair),
52                gas_objects: Arc::new(gas_object_refs),
53            },
54        );
55        genesis_gas_objects.extend(gas_objects);
56    }
57    (accounts, genesis_gas_objects)
58}