1  // Copyright (c) 2025 ADnet Contributors
  2  // This file is part of the AlphaVM library.
  3  
  4  // Licensed under the Apache License, Version 2.0 (the "License");
  5  // you may not use this file except in compliance with the License.
  6  // You may obtain a copy of the License at:
  7  
  8  // http://www.apache.org/licenses/LICENSE-2.0
  9  
 10  // Unless required by applicable law or agreed to in writing, software
 11  // distributed under the License is distributed on an "AS IS" BASIS,
 12  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 13  // See the License for the specific language governing permissions and
 14  // limitations under the License.
 15  
 16  use super::*;
 17  use crate::MIN_VALIDATOR_STAKE;
 18  use console::account::PrivateKey;
 19  
 20  use anyhow::Result;
 21  use proptest::{
 22      collection::{SizeRange, hash_set},
 23      prelude::{Arbitrary, BoxedStrategy, Just, Strategy, any},
 24      sample::size_range,
 25  };
 26  use rand::SeedableRng;
 27  use std::{
 28      collections::HashSet,
 29      hash::{Hash, Hasher},
 30  };
 31  use test_strategy::proptest;
 32  
 33  type CurrentNetwork = console::network::MainnetV0;
 34  
 35  #[derive(Debug, Clone)]
 36  pub struct Validator {
 37      pub private_key: PrivateKey<CurrentNetwork>,
 38      pub address: Address<CurrentNetwork>,
 39      pub stake: u64,
 40      pub is_open: bool,
 41      pub commission: u8,
 42  }
 43  
 44  impl Arbitrary for Validator {
 45      type Parameters = ();
 46      type Strategy = BoxedStrategy<Validator>;
 47  
 48      fn arbitrary_with(_: Self::Parameters) -> Self::Strategy {
 49          any_valid_validator()
 50      }
 51  }
 52  
 53  impl PartialEq<Self> for Validator {
 54      fn eq(&self, other: &Self) -> bool {
 55          self.address == other.address
 56      }
 57  }
 58  
 59  impl Eq for Validator {}
 60  
 61  impl Hash for Validator {
 62      fn hash<H: Hasher>(&self, state: &mut H) {
 63          self.address.hash(state);
 64      }
 65  }
 66  
 67  fn to_committee((round, ValidatorSet(validators)): (u64, ValidatorSet)) -> Result<Committee<CurrentNetwork>> {
 68      Committee::new(round, validators.iter().map(|v| (v.address, (v.stake, v.is_open, v.commission))).collect())
 69  }
 70  
 71  #[derive(Debug, Clone)]
 72  pub struct CommitteeContext(pub Committee<CurrentNetwork>, pub ValidatorSet);
 73  
 74  impl Default for CommitteeContext {
 75      fn default() -> Self {
 76          let validators = ValidatorSet::default();
 77          let committee = to_committee((u64::default(), validators.clone())).unwrap();
 78          Self(committee, validators)
 79      }
 80  }
 81  
 82  impl Arbitrary for CommitteeContext {
 83      type Parameters = ValidatorSet;
 84      type Strategy = BoxedStrategy<CommitteeContext>;
 85  
 86      fn arbitrary() -> Self::Strategy {
 87          any::<ValidatorSet>()
 88              .prop_map(|validators| CommitteeContext(to_committee((1, validators.clone())).unwrap(), validators))
 89              .boxed()
 90      }
 91  
 92      fn arbitrary_with(validator_set: Self::Parameters) -> Self::Strategy {
 93          Just(validator_set)
 94              .prop_map(|validators| CommitteeContext(to_committee((1, validators.clone())).unwrap(), validators))
 95              .boxed()
 96      }
 97  }
 98  
 99  fn validator_set<T: Strategy<Value = Validator>>(
100      element: T,
101      size: impl Into<SizeRange>,
102  ) -> impl Strategy<Value = ValidatorSet> {
103      hash_set(element, size).prop_map(ValidatorSet)
104  }
105  
106  #[derive(Debug, Clone)]
107  pub struct ValidatorSet(pub HashSet<Validator>);
108  
109  impl Default for ValidatorSet {
110      fn default() -> Self {
111          ValidatorSet(
112              (0..4u64)
113                  .map(|i| {
114                      let rng = &mut rand_chacha::ChaChaRng::seed_from_u64(i);
115                      let private_key = PrivateKey::new(rng).unwrap();
116                      let address = Address::try_from(private_key).unwrap();
117                      Validator { private_key, address, stake: MIN_VALIDATOR_STAKE, is_open: false, commission: 0 }
118                  })
119                  .collect(),
120          )
121      }
122  }
123  
124  impl Arbitrary for ValidatorSet {
125      type Parameters = ();
126      type Strategy = BoxedStrategy<ValidatorSet>;
127  
128      fn arbitrary_with(_: Self::Parameters) -> Self::Strategy {
129          // use minimal validator set to speed up tests that require signing from the committee members
130          validator_set(any_valid_validator(), size_range(3..=4usize)).boxed()
131      }
132  }
133  
134  pub fn any_valid_validator() -> BoxedStrategy<Validator> {
135      (MIN_VALIDATOR_STAKE..100_000_000_000_000, any_valid_private_key(), any::<bool>(), 0..100u8)
136          .prop_map(|(stake, private_key, is_open, commission)| {
137              let address = Address::try_from(private_key).unwrap();
138              Validator { private_key, address, stake, is_open, commission }
139          })
140          .boxed()
141  }
142  
143  pub fn any_valid_private_key() -> BoxedStrategy<PrivateKey<CurrentNetwork>> {
144      any::<u64>()
145          .prop_map(|seed| {
146              let rng = &mut rand_chacha::ChaChaRng::seed_from_u64(seed);
147              PrivateKey::new(rng).unwrap()
148          })
149          .boxed()
150  }
151  
152  #[allow(dead_code)]
153  fn too_small_committee() -> BoxedStrategy<Result<Committee<CurrentNetwork>>> {
154      (1u64.., validator_set(any_valid_validator(), 0..3)).prop_map(to_committee).boxed()
155  }
156  
157  #[allow(dead_code)]
158  fn too_low_stake_committee() -> BoxedStrategy<Result<Committee<CurrentNetwork>>> {
159      (1u64.., validator_set(invalid_stake_validator(), 4..=4)).prop_map(to_committee).boxed()
160  }
161  
162  #[allow(dead_code)]
163  fn invalid_stake_validator() -> BoxedStrategy<Validator> {
164      (0..MIN_VALIDATOR_STAKE, any_valid_private_key(), any::<bool>(), 0..u8::MAX)
165          .prop_map(|(stake, private_key, is_open, commission)| {
166              let address = Address::try_from(private_key).unwrap();
167              Validator { private_key, address, stake, is_open, commission }
168          })
169          .boxed()
170  }
171  
172  #[proptest]
173  fn committee_members(input: CommitteeContext) {
174      let CommitteeContext(committee, ValidatorSet(validators)) = input;
175  
176      let mut total_stake = 0u64;
177      for v in validators.iter() {
178          total_stake += v.stake;
179      }
180  
181      assert_eq!(committee.num_members(), validators.len());
182      assert_eq!(committee.total_stake(), total_stake);
183      for v in validators.iter() {
184          let address = v.address;
185          assert!(committee.is_committee_member(address));
186          assert_eq!(committee.get_stake(address), v.stake);
187      }
188      let quorum_threshold = committee.quorum_threshold();
189      let availability_threshold = committee.availability_threshold();
190      // (N - f) + (f + 1) - 1 = N
191      assert_eq!(quorum_threshold + availability_threshold - 1, total_stake);
192  }
193  
194  #[proptest]
195  fn invalid_stakes(#[strategy(too_low_stake_committee())] committee: Result<Committee<CurrentNetwork>>) {
196      assert!(committee.is_err());
197      if let Err(err) = committee {
198          assert_eq!(err.to_string().as_str(), "All members must have at least 10000000000000 microcredits in stake");
199      }
200  }
201  
202  #[proptest]
203  fn invalid_member_count(#[strategy(too_small_committee())] committee: Result<Committee<CurrentNetwork>>) {
204      assert!(matches!(committee, Err(e) if e.to_string().as_str() == "Committee must have at least 3 members"))
205  }