unit tests for controlled module (#428)

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

@@ -6,7 +6,7 @@ use crate::{
 };
 use chain_addr::{Address, Kind};
 use imhamt::Hamt;
-use std::{collections::hash_map::DefaultHasher, num::NonZeroU64};
+use std::{collections::hash_map::DefaultHasher, fmt, num::NonZeroU64};
 
 #[derive(Default, Clone, Eq, PartialEq)]
 pub struct StakeControl {
@@ -158,3 +158,253 @@ impl StakeControl {
         }
     }
 }
+
+impl fmt::Debug for StakeControl {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(
+            f,
+            "unassigned: {}, assigned: {}, control: {:?}",
+            self.unassigned,
+            self.assigned,
+            self.control
+                .iter()
+                .map(|(id, account)| (id.clone(), account.clone()))
+                .collect::<Vec<(Identifier, Stake)>>()
+        )
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::StakeControl;
+    use crate::{
+        account::{self, Identifier},
+        rewards::Ratio,
+        stake::Stake,
+        testing::{utxo::ArbitaryLedgerUtxo, TestGen},
+    };
+    use quickcheck_macros::quickcheck;
+    use std::num::NonZeroU64;
+
+    fn create_stake_control_from(
+        assigned: &[(Identifier, Stake)],
+        unassigned: Stake,
+    ) -> StakeControl {
+        let stake_control: StakeControl = assigned
+            .iter()
+            .fold(StakeControl::new(), |sc, (identifier, stake)| {
+                sc.add_to(identifier.clone(), *stake)
+            });
+        stake_control.add_unassigned(unassigned)
+    }
+
+    #[test]
+    pub fn empty_stake_control() {
+        let random_identifier = TestGen::identifier();
+        let stake_control = create_stake_control_from(&[], Stake::zero());
+
+        assert_eq!(stake_control.total(), Stake::zero());
+        assert_eq!(stake_control.unassigned(), Stake::zero());
+        assert_eq!(stake_control.assigned(), Stake::zero());
+        assert_eq!(stake_control.by(&random_identifier), None);
+        let expected_ratio = Ratio {
+            numerator: 0,
+            denominator: NonZeroU64::new(1).unwrap(),
+        };
+        assert_eq!(stake_control.ratio_by(&random_identifier), expected_ratio);
+    }
+
+    #[test]
+    pub fn stake_control_only_assigned() {
+        let identifier = TestGen::identifier();
+        let initial_stake = Stake(100);
+        let stake_control =
+            create_stake_control_from(&[(identifier.clone(), initial_stake)], Stake::zero());
+
+        assert_eq!(stake_control.total(), initial_stake);
+        assert_eq!(stake_control.unassigned(), Stake::zero());
+        assert_eq!(stake_control.assigned(), initial_stake);
+        assert_eq!(stake_control.by(&identifier).unwrap(), initial_stake);
+        let expected_ratio = Ratio {
+            numerator: 100,
+            denominator: NonZeroU64::new(100).unwrap(),
+        };
+        assert_eq!(stake_control.ratio_by(&identifier), expected_ratio);
+    }
+
+    #[test]
+    pub fn stake_control_only_unassigned() {
+        let identifier = TestGen::identifier();
+        let initial_stake = Stake(100);
+        let stake_control = create_stake_control_from(&[], initial_stake);
+
+        assert_eq!(stake_control.total(), initial_stake);
+        assert_eq!(stake_control.unassigned(), initial_stake);
+        assert_eq!(stake_control.assigned(), Stake::zero());
+        assert_eq!(stake_control.by(&identifier), None);
+        let expected_ratio = Ratio {
+            numerator: 0,
+            denominator: NonZeroU64::new(1).unwrap(),
+        };
+        assert_eq!(stake_control.ratio_by(&identifier), expected_ratio);
+    }
+
+    #[test]
+    pub fn stake_control_unassigned_and_assigned() {
+        let identifier = TestGen::identifier();
+        let stake_to_add = Stake(100);
+
+        let stake_control =
+            create_stake_control_from(&[(identifier.clone(), stake_to_add.clone())], stake_to_add);
+
+        assert_eq!(stake_control.total(), Stake(200));
+        assert_eq!(stake_control.unassigned(), stake_to_add);
+        assert_eq!(stake_control.assigned(), stake_to_add);
+        assert_eq!(stake_control.by(&identifier), Some(stake_to_add));
+        let expected_ratio = Ratio {
+            numerator: 100,
+            denominator: NonZeroU64::new(100).unwrap(),
+        };
+        assert_eq!(stake_control.ratio_by(&identifier), expected_ratio);
+    }
+
+    #[test]
+    pub fn stake_control_remove_part_of_assigned() {
+        let identifier = TestGen::identifier();
+        let stake_to_add = Stake(100);
+        let stake_to_sub = Stake(50);
+        let mut stake_control =
+            create_stake_control_from(&[(identifier.clone(), stake_to_add.clone())], stake_to_add);
+        stake_control = stake_control.remove_from(identifier.clone(), stake_to_sub.clone());
+
+        assert_eq!(stake_control.total(), Stake(150));
+        assert_eq!(stake_control.unassigned(), Stake(100));
+        assert_eq!(stake_control.assigned(), Stake(50));
+        assert_eq!(stake_control.by(&identifier), Some(Stake(50)));
+        let expected_ratio = Ratio {
+            numerator: 50,
+            denominator: NonZeroU64::new(50).unwrap(),
+        };
+        assert_eq!(stake_control.ratio_by(&identifier), expected_ratio);
+    }
+
+    #[test]
+    #[should_panic(expected = "KeyNotFound")]
+    pub fn stake_control_remove_non_exsiting_assigned() {
+        let non_existing_identifier = TestGen::identifier();
+        let existing_identifier = TestGen::identifier();
+        let stake_to_add = Stake(100);
+
+        let stake_control =
+            create_stake_control_from(&[(existing_identifier, stake_to_add.clone())], stake_to_add);
+
+        assert_eq!(stake_control.total(), Stake(200));
+        let _ = stake_control.remove_from(non_existing_identifier, stake_to_add);
+    }
+
+    #[test]
+    pub fn stake_control_remove_all_assigned() {
+        let identifier = TestGen::identifier();
+        let stake_to_add = Stake(100);
+
+        let mut stake_control =
+            create_stake_control_from(&[(identifier.clone(), stake_to_add.clone())], stake_to_add);
+
+        assert_eq!(stake_control.total(), Stake(200));
+
+        stake_control = stake_control.remove_from(identifier.clone(), stake_to_add);
+
+        assert_eq!(stake_control.total(), Stake(100));
+        assert_eq!(stake_control.unassigned(), Stake(100));
+        assert_eq!(stake_control.assigned(), Stake::zero());
+        assert_eq!(stake_control.by(&identifier), Some(Stake::zero()));
+        unsafe {
+            let expected_ratio = Ratio {
+                numerator: 0,
+                denominator: NonZeroU64::new_unchecked(0),
+            };
+            assert_eq!(stake_control.ratio_by(&identifier), expected_ratio);
+        }
+    }
+
+    #[test]
+    pub fn stake_control_remove_unassigned() {
+        let identifier = TestGen::identifier();
+        let stake_to_add = Stake(100);
+
+        let stake_control =
+            create_stake_control_from(&[(identifier.clone(), stake_to_add.clone())], stake_to_add);
+
+        assert_eq!(stake_control.total(), Stake(200));
+        assert_eq!(stake_control.unassigned(), stake_to_add);
+        assert_eq!(stake_control.assigned(), stake_to_add);
+        assert_eq!(stake_control.by(&identifier), Some(stake_to_add));
+        let expected_ratio = Ratio {
+            numerator: 100,
+            denominator: NonZeroU64::new(100).unwrap(),
+        };
+        assert_eq!(stake_control.ratio_by(&identifier), expected_ratio);
+    }
+
+    #[test]
+    pub fn stake_control_remove_all() {
+        let identifier = TestGen::identifier();
+        let stake_to_add = Stake(100);
+
+        let mut stake_control =
+            create_stake_control_from(&[(identifier.clone(), stake_to_add.clone())], stake_to_add);
+
+        stake_control = stake_control.remove_from(identifier.clone(), stake_to_add.clone());
+        stake_control = stake_control.remove_unassigned(stake_to_add.clone());
+
+        assert_eq!(stake_control.total(), Stake::zero());
+        assert_eq!(stake_control.unassigned(), Stake::zero());
+        assert_eq!(stake_control.assigned(), Stake::zero());
+        assert_eq!(stake_control.by(&identifier), Some(Stake::zero()));
+    }
+
+    #[test]
+    pub fn stake_control_account_ratio() {
+        let first_identifier = TestGen::identifier();
+        let second_identifier = TestGen::identifier();
+        let stake_to_add = Stake(100);
+
+        let stake_control = create_stake_control_from(
+            &[
+                (first_identifier.clone(), stake_to_add.clone()),
+                (second_identifier.clone(), stake_to_add.clone()),
+            ],
+            stake_to_add,
+        );
+
+        assert_eq!(stake_control.by(&first_identifier), Some(stake_to_add));
+        assert_eq!(stake_control.by(&second_identifier), Some(stake_to_add));
+
+        assert_eq!(stake_control.by(&first_identifier), Some(stake_to_add));
+        assert_eq!(stake_control.by(&second_identifier), Some(stake_to_add));
+
+        let expected_ratio = Ratio {
+            numerator: 100,
+            denominator: NonZeroU64::new(200).unwrap(),
+        };
+        assert_eq!(stake_control.ratio_by(&first_identifier), expected_ratio);
+
+        let expected_ratio = Ratio {
+            numerator: 100,
+            denominator: NonZeroU64::new(200).unwrap(),
+        };
+        assert_eq!(stake_control.ratio_by(&second_identifier), expected_ratio);
+    }
+
+    #[quickcheck]
+    pub fn stake_control_from_ledger(accounts: account::Ledger, utxos: ArbitaryLedgerUtxo) {
+        let stake_control = StakeControl::new_with(&accounts, &utxos.0);
+        //verify sum
+        let accounts = accounts.get_total_value().unwrap();
+        let utxo_or_group = utxos.0.values().map(|x| x.value).sum();
+        let expected_sum = accounts
+            .checked_add(utxo_or_group)
+            .expect("cannot calculate expected total");
+        assert_eq!(stake_control.total(), expected_sum.into());
+    }
+}

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

@@ -6,6 +6,7 @@ pub mod output;
 pub mod transaction;
 pub mod update_proposal;
 pub mod utils;
+pub mod utxo;
 pub mod wallet;
 
 use crate::{key::BftLeaderId, transaction::Output, value::Value};

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

@@ -0,0 +1,38 @@
+use super::AverageValue;
+use crate::{key::Hash, testing::data::AddressData, transaction::Output, utxo::Ledger};
+use chain_addr::{Address, Discrimination};
+use quickcheck::{Arbitrary, Gen};
+use std::{collections::HashMap, iter};
+
+#[derive(Debug, Clone)]
+pub struct ArbitaryLedgerUtxo(pub Ledger<Address>);
+
+impl Arbitrary for ArbitaryLedgerUtxo {
+    fn arbitrary<G: Gen>(g: &mut G) -> Self {
+        let mut ledger = Ledger::new();
+        let size = usize::arbitrary(g) % 50 + 1;
+        let arbitrary_utxos: HashMap<Hash, (u8, Output<Address>)> = iter::from_fn(|| {
+            let outs = match u8::arbitrary(g) % 2 {
+                0 => (
+                    0u8,
+                    AddressData::utxo(Discrimination::Test)
+                        .make_output(AverageValue::arbitrary(g).into()),
+                ),
+                1 => (
+                    0u8,
+                    AddressData::delegation(Discrimination::Test)
+                        .make_output(AverageValue::arbitrary(g).into()),
+                ),
+                _ => unreachable!(),
+            };
+            Some((Hash::arbitrary(g), outs))
+        })
+        .take(size)
+        .collect();
+
+        for (key, value) in arbitrary_utxos {
+            ledger = ledger.add(&key, &[value]).unwrap();
+        }
+        ArbitaryLedgerUtxo(ledger)
+    }
+}