1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
use crate::Wallet;
use jortestkit::load::{Request, RequestFailure, RequestGenerator};
use rand::Rng;
use rand_core::OsRng;
use std::time::Instant;
use valgrind::WalletNodeRestClient;

pub struct AccountRequestGen {
    rand: OsRng,
    wallets: Vec<Wallet>,
    rest: WalletNodeRestClient,
}

impl AccountRequestGen {
    pub fn new(wallets: Vec<Wallet>, rest: WalletNodeRestClient) -> Self {
        Self {
            wallets,
            rand: OsRng,
            rest,
        }
    }

    pub fn random_account_request(&mut self) -> Result<(), RequestFailure> {
        let wallet_index = self.rand.gen_range(0..self.wallets.len());
        let wallet = self.wallets.get(wallet_index).ok_or_else(|| {
            RequestFailure::General(format!("wallet with index '{wallet_index}' not found"))
        })?;
        self.rest
            .account_state(wallet.id())
            .map_err(|e| RequestFailure::General(format!("{e:?}")))?;
        Ok(())
    }
}

impl RequestGenerator for AccountRequestGen {
    fn next(&mut self) -> Result<Request, RequestFailure> {
        let start = Instant::now();
        match self.random_account_request() {
            Ok(()) => Ok(Request {
                ids: vec![],
                duration: start.elapsed(),
            }),
            Err(e) => Err(RequestFailure::General(format!("{e:?}"))),
        }
    }

    fn split(mut self) -> (Self, Option<Self>) {
        let wallets_len = self.wallets.len();
        if wallets_len <= 1 {
            return (self, None);
        }
        let wallets = self.wallets.split_off(wallets_len / 2);
        let new_gen = Self {
            rand: self.rand,
            wallets,
            rest: self.rest.clone(),
        };

        (self, Some(new_gen))
    }
}