Property tests for stake distribution.

Signed-off-by: dkijania <[email protected]>

chain-impl-mockchain/src/stake/distribution.rs

@@ -129,3 +129,285 @@ pub fn get_distribution(
         to_pools: dist,
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use crate::account;
+    use crate::stake::delegation::DelegationState;
+    use crate::{
+        account::{AccountAlg, Identifier},
+        fragment::FragmentId,
+        stake::StakePoolInfo,
+        testing::{
+            address::{AddressData, AddressDataValue},
+            arbitrary::utils as arbitrary_utils,
+            arbitrary::ArbitraryAddressDataValueVec,
+        },
+        transaction::{Output, TransactionIndex},
+        utxo,
+        value::Value,
+    };
+    use chain_addr::{Address, Kind};
+    use chain_crypto::PublicKey;
+    use quickcheck::{Arbitrary, Gen, TestResult};
+    use quickcheck_macros::quickcheck;
+
+    /// Holds all possible cases of distribution source
+    #[derive(Clone, Debug)]
+    pub struct StakeDistributionArbitraryData {
+        utxos: Vec<(FragmentId, TransactionIndex, Output<Address>)>,
+        unassigned_accounts: Vec<(Identifier, Value)>,
+        assigned_accounts: Vec<(Identifier, Value)>,
+        dangling_accounts: Vec<(Identifier, Value)>,
+        groups: Vec<(FragmentId, TransactionIndex, Output<Address>)>,
+        groups_single_account: Vec<(FragmentId, TransactionIndex, Output<Address>)>,
+        single_account: (Identifier, Value),
+        active_stake_pool: StakePoolInfo,
+        retired_stake_pool: StakePoolInfo,
+    }
+
+    impl Arbitrary for StakeDistributionArbitraryData {
+        fn arbitrary<G: Gen>(gen: &mut G) -> Self {
+            let arbitrary_addresses = ArbitraryAddressDataValueVec::arbitrary(gen);
+
+            let utxos = arbitrary_addresses
+                .utxos()
+                .iter()
+                .map(|x| {
+                    (
+                        FragmentId::arbitrary(gen),
+                        TransactionIndex::arbitrary(gen),
+                        x.make_output(),
+                    )
+                })
+                .collect();
+
+            let splitted_vec_of_accounts =
+                arbitrary_utils::split_vec(&arbitrary_addresses.accounts(), gen, 3);
+
+            let unassigned_accounts = splitted_vec_of_accounts
+                .get(0)
+                .unwrap()
+                .iter()
+                .map(|x| (Identifier::from(x.address_data.public_key()), x.value))
+                .collect();
+
+            let assigned_accounts = splitted_vec_of_accounts
+                .get(1)
+                .unwrap()
+                .iter()
+                .map(|x| (Identifier::from(x.address_data.public_key()), x.value))
+                .collect();
+
+            let dangling_accounts = splitted_vec_of_accounts
+                .get(2)
+                .unwrap()
+                .iter()
+                .map(|x| (Identifier::from(x.address_data.public_key()), x.value))
+                .collect();
+
+            let splitted_vec_of_delegations =
+                arbitrary_utils::split_vec(&arbitrary_addresses.delegations(), gen, 2);
+
+            let groups = splitted_vec_of_delegations
+                .get(0)
+                .unwrap()
+                .iter()
+                .map(|x| {
+                    (
+                        FragmentId::arbitrary(gen),
+                        TransactionIndex::arbitrary(gen),
+                        x.make_output(),
+                    )
+                })
+                .collect();
+
+            let single_account = (Identifier::arbitrary(gen), Value::arbitrary(gen));
+
+            let groups_single_account = splitted_vec_of_delegations
+                .get(1)
+                .unwrap()
+                .iter()
+                .map(|x| {
+                    AddressDataValue::new(
+                        AddressData::delegation_for_account(
+                            x.address_data.clone(),
+                            single_account.0.clone().into(),
+                        ),
+                        x.value.clone(),
+                    )
+                })
+                .map(|x| {
+                    (
+                        FragmentId::arbitrary(gen),
+                        TransactionIndex::arbitrary(gen),
+                        x.make_output(),
+                    )
+                })
+                .collect();
+
+            let active_stake_pool = StakePoolInfo::arbitrary(gen);
+            let retired_stake_pool = StakePoolInfo::arbitrary(gen);
+
+            StakeDistributionArbitraryData {
+                utxos,
+                unassigned_accounts,
+                assigned_accounts,
+                dangling_accounts,
+                groups,
+                groups_single_account,
+                single_account,
+                active_stake_pool,
+                retired_stake_pool,
+            }
+        }
+    }
+
+    impl StakeDistributionArbitraryData {
+        pub fn calculate_unassigned(&self) -> Value {
+            Value(
+                self.get_sum_from_utxo_type(&self.utxos)
+                    + self.get_sum_from_account_type(&self.unassigned_accounts),
+            )
+        }
+
+        pub fn calculate_dangling(&self) -> Value {
+            Value(self.get_sum_from_account_type(&self.dangling_accounts))
+        }
+
+        pub fn pools_total(&self) -> u64 {
+            self.get_sum_from_account_type(&self.assigned_accounts)
+                + self.get_sum_from_utxo_type(&self.groups)
+                + self.get_sum_from_utxo_type(&self.groups_single_account)
+                + (&self.single_account.1).0
+        }
+
+        fn get_sum_from_utxo_type(
+            &self,
+            utxos: &Vec<(FragmentId, TransactionIndex, Output<Address>)>,
+        ) -> u64 {
+            utxos.iter().map(|(_, _, x)| x.value.0).sum::<u64>()
+        }
+
+        fn get_sum_from_account_type(&self, accounts: &Vec<(Identifier, Value)>) -> u64 {
+            accounts.iter().map(|(_, x)| x.0).sum::<u64>()
+        }
+    }
+
+    #[quickcheck]
+    pub fn stake_distribution_is_consistent_with_total_value(
+        stake_distribution_data: StakeDistributionArbitraryData,
+    ) -> TestResult {
+        let mut accounts = account::Ledger::new();
+        let mut dstate = DelegationState::new();
+        let mut utxos = utxo::Ledger::new();
+
+        // create two stake pools, one active and one inactive
+        let id_active_pool = stake_distribution_data.active_stake_pool.to_id();
+        dstate = dstate
+            .register_stake_pool(stake_distribution_data.active_stake_pool.clone())
+            .unwrap();
+        let id_retired_pool = stake_distribution_data.retired_stake_pool.to_id();
+
+        // add utxos
+        for (fragment_id, tx_index, output) in stake_distribution_data.utxos.iter().cloned() {
+            utxos = utxos.add(&fragment_id, &[(tx_index, output)]).unwrap();
+        }
+
+        // add delegation addresses with all accounts delegated to active stake pool
+        for (fragment_id, tx_index, output) in stake_distribution_data.groups.iter().cloned() {
+            utxos = utxos
+                .add(&fragment_id, &[(tx_index, output.clone())])
+                .unwrap();
+            let account_public_key: PublicKey<AccountAlg> = match output.address.kind() {
+                Kind::Group(_, delegation_key) => delegation_key.clone(),
+                _ => panic!("delegation utxo should have Kind::Group type"),
+            };
+            accounts = accounts
+                .add_account(
+                    &Identifier::from(account_public_key.clone()),
+                    Value::zero(),
+                    (),
+                )
+                .unwrap();
+            accounts = accounts
+                .set_delegation(
+                    &Identifier::from(account_public_key.clone()),
+                    Some(id_active_pool.clone()),
+                )
+                .unwrap();
+        }
+
+        // add delegation addresses which point to single account with delegation
+        for (fragment_id, tx_index, output) in stake_distribution_data
+            .groups_single_account
+            .iter()
+            .cloned()
+        {
+            utxos = utxos.add(&fragment_id, &[(tx_index, output)]).unwrap();
+        }
+
+        // add accounts without delegation
+        for (id, value) in stake_distribution_data.unassigned_accounts.iter().cloned() {
+            accounts = accounts.add_account(&id, value, ()).unwrap();
+        }
+
+        // add accounts with delegation
+        for (id, value) in stake_distribution_data.assigned_accounts.iter().cloned() {
+            accounts = accounts.add_account(&id, value, ()).unwrap();
+            accounts = accounts
+                .set_delegation(&id, Some(id_active_pool.clone()))
+                .unwrap();
+        }
+
+        // add accounts with delegation as a target for delegation addresses
+        let single_account = stake_distribution_data.single_account.clone();
+        accounts = accounts
+            .add_account(&single_account.0.clone(), single_account.1, ())
+            .unwrap();
+        accounts = accounts
+            .set_delegation(&single_account.0.clone(), Some(id_active_pool.clone()))
+            .unwrap();
+
+        // add accounts with retired stake pool
+        for (id, value) in stake_distribution_data.dangling_accounts.iter().cloned() {
+            accounts = accounts.add_account(&id, value, ()).unwrap();
+            accounts = accounts
+                .set_delegation(&id, Some(id_retired_pool.clone()))
+                .unwrap();
+        }
+
+        // verify
+        let distribution = super::get_distribution(&accounts, &dstate, &utxos);
+
+        if distribution.unassigned != stake_distribution_data.calculate_unassigned() {
+            return TestResult::error(format!(
+                "Wrong Unassigned value. expected: {} but got {}",
+                stake_distribution_data.calculate_unassigned(),
+                &distribution.unassigned
+            ));
+        }
+
+        if distribution.dangling != stake_distribution_data.calculate_dangling() {
+            return TestResult::error(format!(
+                "Wrong Unassigned value. expected: {} but got {}",
+                stake_distribution_data.calculate_unassigned(),
+                &distribution.unassigned
+            ));
+        }
+
+        let pools_total_stake = distribution
+            .to_pools
+            .values()
+            .map(|x| x.total_stake.0)
+            .sum::<u64>();
+        if pools_total_stake != stake_distribution_data.pools_total() {
+            return TestResult::error(format!(
+                "Wrong Unassigned value. expected: {} but got {}",
+                stake_distribution_data.pools_total(),
+                pools_total_stake
+            ));
+        }
+        TestResult::passed()
+    }
+}

chain-impl-mockchain/src/testing/address.rs

@@ -54,7 +54,6 @@ impl AddressData {
         }
     }
 
-
     pub fn from_discrimination_and_kind_type(
         discrimination: Discrimination,
         kind: &KindType,
@@ -152,6 +151,21 @@ impl AddressData {
     pub fn generate_key_pair<A: AsymmetricKey>() -> KeyPair<A> {
         TestCryptoGen(0).keypair::<A>(rand_os::OsRng::new().unwrap().next_u32())
     }
+
+    pub fn delegation_for_account(
+        other: AddressData,
+        delegation_public_key: PublicKey<Ed25519>,
+    ) -> Self {
+        let user_address = Address(
+            other.address.discrimination().clone(),
+            Kind::Group(other.public_key().clone(), delegation_public_key.clone()),
+        );
+        AddressData::new(other.private_key, other.spending_counter, user_address)
+    }
+
+    fn generate_random_secret_key() -> EitherEd25519SecretKey {
+        EitherEd25519SecretKey::generate(rand_os::OsRng::new().unwrap())
+    }
 }
 
 #[derive(Clone, Debug, PartialEq)]

chain-impl-mockchain/src/testing/arbitrary/address.rs

@@ -1,5 +1,9 @@
-use crate::testing::{address::AddressData, arbitrary::kind_type::KindTypeWithoutMultisig};
-use chain_addr::Discrimination;
+use crate::testing::{
+    address::{AddressData, AddressDataValue},
+    arbitrary::kind_type::KindTypeWithoutMultisig,
+    arbitrary::AverageValue,
+};
+use chain_addr::{Discrimination, Kind};
 use quickcheck::{Arbitrary, Gen};
 use std::iter;
 
@@ -24,3 +28,58 @@ impl Arbitrary for AddressData {
         )
     }
 }
+
+impl Arbitrary for AddressDataValue {
+    fn arbitrary<G: Gen>(gen: &mut G) -> Self {
+        AddressDataValue::new(
+            Arbitrary::arbitrary(gen),
+            AverageValue::arbitrary(gen).into(),
+        )
+    }
+}
+
+#[derive(Clone, Debug)]
+pub struct ArbitraryAddressDataValueVec(pub Vec<AddressDataValue>);
+
+impl Arbitrary for ArbitraryAddressDataValueVec {
+    fn arbitrary<G: Gen>(gen: &mut G) -> Self {
+        let size_limit = 10;
+        let n = usize::arbitrary(gen) % size_limit + 1;
+        let addresses = iter::from_fn(|| Some(AddressDataValue::arbitrary(gen))).take(n);
+        ArbitraryAddressDataValueVec(addresses.collect())
+    }
+}
+
+impl ArbitraryAddressDataValueVec {
+    pub fn utxos(&self) -> Vec<AddressDataValue> {
+        self.0
+            .iter()
+            .cloned()
+            .filter(|x| match x.address_data.kind() {
+                Kind::Single { .. } => true,
+                _ => false,
+            })
+            .collect()
+    }
+    pub fn accounts(&self) -> Vec<AddressDataValue> {
+        self.0
+            .iter()
+            .cloned()
+            .filter(|x| match x.address_data.kind() {
+                Kind::Account { .. } => true,
+                _ => false,
+            })
+            .collect()
+    }
+
+    pub fn delegations(&self) -> Vec<AddressDataValue> {
+        self.0
+            .iter()
+            .cloned()
+            .filter(|x| match x.address_data.kind() {
+                Kind::Group { .. } => true,
+                _ => false,
+            })
+            .collect()
+    }
+}