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
62
63
64
65
66
67
68
69
70

use crate::{signed_delegation_cert, signed_stake_pool_cert, StakePool, Wallet};
use jormungandr_automation::jormungandr::Block0ConfigurationBuilder;
use jormungandr_lib::interfaces::{Initial, InitialUTxO, Value};

pub trait Block0ConfigurationBuilderExtension {
    fn with_wallet(self, wallet: &Wallet, value: Value) -> Self;
    fn with_wallets_having_some_values(self, wallets: Vec<&Wallet>) -> Self;
    fn with_stake_pool(self, stake_pool: &StakePool) -> Self;
    fn with_delegation_to_stake_pool(self, stake_pool: &StakePool, wallets: Vec<&Wallet>) -> Self;
    fn with_stake_pool_and_delegation(self, stake_pool: &StakePool, wallets: Vec<&Wallet>) -> Self;
}

impl Block0ConfigurationBuilderExtension for Block0ConfigurationBuilder {
    fn with_wallet(self, wallet: &Wallet, value: Value) -> Self {
        self.with_funds(vec![Initial::Fund(vec![InitialUTxO {
            address: wallet.address(),
            value,
        }])])
    }

    fn with_wallets_having_some_values(mut self, wallets: Vec<&Wallet>) -> Self {
        for wallet in wallets {
            self = self.with_wallet(wallet, 1_000_000.into());
        }
        self
    }

    fn with_stake_pool(self, stake_pool: &StakePool) -> Self {
        self.with_certs(vec![Initial::Cert(
            signed_stake_pool_cert(
                chain_impl_mockchain::block::BlockDate {
                    epoch: 1,
                    slot_id: 0,
                },
                stake_pool,
            )
            .into(),
        )])
    }

    fn with_delegation_to_stake_pool(self, stake_pool: &StakePool, wallets: Vec<&Wallet>) -> Self {
        self.with_certs(
            wallets
                .iter()
                .map(|x| {
                    Initial::Cert(
                        signed_delegation_cert(
                            x,
                            chain_impl_mockchain::block::BlockDate {
                                epoch: 1,
                                slot_id: 0,
                            },
                            stake_pool.inner().id(),
                        )
                        .into(),
                    )
                })
                .collect(),
        )
    }

    fn with_stake_pool_and_delegation(
        self,
        stake_pool: &StakePool,
        delegators: Vec<&Wallet>,
    ) -> Self {
        self.with_stake_pool(stake_pool)
            .with_delegation_to_stake_pool(stake_pool, delegators)
    }
}