hotshot_testing/
node_stake.rs

1use std::collections::HashMap;
2
3use alloy::primitives::U256;
4
5#[derive(Clone)]
6pub struct TestNodeStakes {
7    stakes: HashMap<u64, U256>,
8    default_stake: U256,
9}
10
11impl TestNodeStakes {
12    pub fn new(stakes: HashMap<u64, U256>, default_stake: U256) -> Self {
13        Self {
14            stakes,
15            default_stake,
16        }
17    }
18
19    pub fn get(&self, node_id: u64) -> U256 {
20        self.stakes
21            .get(&node_id)
22            .cloned()
23            .unwrap_or(self.default_stake)
24    }
25
26    pub fn with_stake(mut self, node_id: u64, stake: U256) -> Self {
27        self.stakes.insert(node_id, stake);
28        self
29    }
30}
31
32impl Default for TestNodeStakes {
33    fn default() -> Self {
34        Self {
35            stakes: HashMap::new(),
36            default_stake: U256::from(1),
37        }
38    }
39}
hotshot_testing/node_stake.rs

hotshot_testing/
node_stake.rs
