Add deterministic leader schedule test (#9096)

ledger/src/leader_schedule.rs

@@ -113,7 +113,7 @@ fn sort_stakes(stakes: &mut Vec<(&Pubkey, u64)>) {
 
 #[cfg(test)]
 mod tests {
-    use {super::*, itertools::Itertools, rand::Rng, std::iter::repeat_with};
+    use {super::*, itertools::Itertools, rand::Rng, std::iter::repeat_with, test_case::test_case};
 
     #[test]
     fn test_get_leader_upcoming_slots() {
@@ -166,4 +166,102 @@ mod tests {
         sort_stakes(&mut stakes);
         assert_eq!(stakes, vec![(&pubkey1, 1), (&pubkey0, 1)]);
     }
+
+    fn pubkey_from_u16(n: u16) -> Pubkey {
+        let mut bytes = [0; 32];
+        bytes[0..2].copy_from_slice(&n.to_le_bytes());
+        Pubkey::new_from_array(bytes)
+    }
+
+    #[test_case(1, &[10, 20, 30], 12, 1, &[1, 1, 2, 1, 1, 0, 0, 1, 2, 1, 0, 1])]
+    #[test_case(1, &[10, 20, 30], 12, 2, &[1, 1, 1, 1, 2, 2, 1, 1, 1, 1, 0, 0])]
+    #[test_case(1, &[30, 10, 20], 12, 1, &[2, 2, 0, 2, 2, 1, 1, 2, 0, 2, 1, 2])]
+    #[test_case(1, &[30, 10, 20], 12, 2, &[2, 2, 2, 2, 0, 0, 2, 2, 2, 2, 1,1])]
+    #[test_case(1, &[10, 20, 25, 30], 12, 1, &[2, 2, 3, 1, 2, 0, 1, 1, 3, 2, 1, 2])]
+    #[test_case(1, &[10, 20, 25, 30, 35, 40, 100], 15, 1,
+                &[4, 5, 6, 3, 4, 1, 2, 3, 6, 4, 2, 4, 5, 6, 6])]
+    #[test_case(1, &[10, 20, 25, 30, 35, 40, 100, 1000], 15, 1,
+                &[7, 7, 7, 7, 7, 4, 6, 7, 7, 7, 6, 7, 7, 7, 7])]
+    #[test_case(1, &[10, 20, 25, 30, 35, 40, 100, 1000, 10_000], 20, 1,
+                &[8, 8, 8, 8, 8, 7, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 7])]
+    #[test_case(1, &[10, 20, 25, 30, 35, 40, 100, 1000, 10_000], 25, 1,
+                &[8, 8, 8, 8, 8, 7, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 7, 8, 8, 8, 8, 8])]
+    #[test_case(457468, &[10, 20, 30], 12, 1, &[2, 2, 0, 1, 0, 2, 1, 2, 1, 2, 2, 2])]
+    #[test_case(457468, &[10, 20, 30], 12, 2, &[2, 2, 2, 2, 0, 0, 1, 1, 0, 0, 2, 2])]
+    #[test_case(457469, &[10, 20, 30], 12, 1, &[1, 2, 2, 2, 2, 2, 2, 1, 0, 2, 2, 0])]
+    #[test_case(457470, &[10, 20, 30], 12, 1, &[2, 1, 1, 1, 1, 1, 1, 1, 1, 2, 0, 2])]
+    #[test_case(3466545, &[10, 20, 30], 12, 1, &[2, 2, 0, 0, 2, 1, 1, 1, 0, 0, 2, 2])]
+    #[test_case(3466545, &[10, 20, 30], 13, 1, &[2, 2, 0, 0, 2, 1, 1, 1, 0, 0, 2, 2, 1])]
+    #[test_case(3466545, &[10, 20, 30], 13, 2, &[2, 2, 2, 2, 0, 0, 0, 0, 2, 2, 1, 1, 1])]
+    #[test_case(3466545, &[10, 20, 30], 14, 1, &[2, 2, 0, 0, 2, 1, 1, 1, 0, 0, 2, 2, 1, 2])]
+    #[test_case(3466545, &[10, 20, 30], 14, 2, &[2, 2, 2, 2, 0, 0, 0, 0, 2, 2, 1, 1, 1, 1])]
+    fn test_stake_leader_schedule_exact_order(
+        epoch: u64,
+        stakes: &[u64],
+        len: u64,
+        repeat: u64,
+        expected_order: &[usize],
+    ) {
+        let pubkeys: Vec<_> = (0..stakes.len() as u16).map(pubkey_from_u16).collect();
+        let stakes = pubkeys.iter().zip(stakes.iter().copied()).collect();
+        let order: Vec<_> = stake_weighted_slot_leaders(stakes, epoch, len, repeat)
+            .into_iter()
+            .map(|pubkey| {
+                pubkeys
+                    .iter()
+                    .find_position(|item| *item == &pubkey)
+                    .unwrap()
+                    .0
+            })
+            .collect();
+        assert_eq!(order, expected_order);
+    }
+
+    #[test_case(42, 1_000, 0, "4XU6LEarBUmBkAvXRsjeyLu3N8CcgrvbRFrNiJi2jECk")]
+    #[test_case(42, 10_000, 0, "G2MGFXgdLATXWr1336i8PTcaUMc4GbJRMJdbxiarCttr")]
+    #[test_case(42, 10_000, 1, "9xLLKyyqF5YrdwPSDbqh5oVamSF7cqPqQLxEyHTexEiP")]
+    #[test_case(42, 10_000, 2, "AJ6NQi2p5SnRz9mqESqkW2PwVoT2vYy1fmKdaHxNFUAf")]
+    #[test_case(42, 10_000, 3, "2oLjZggMwDTQhzdB4KN5VQisyeRw6MZbBBdjosNZK5xR")]
+    #[test_case(346436, 1_000, 0, "59SnXMS4NzTSib8TNykiJgFQBeAVxUqsAvQm7JtkodPQ")]
+    #[test_case(346436, 1_000, 1, "BEB2nC9MBALPbgwGKGfHu6V88QG7doScx65cAd6VjnRk")]
+    #[test_case(346436, 1_000, 2, "3aLE5S6xLEU9yg5EZQH27qrC86aC2dG8KLh4NbcapXpy")]
+    #[test_case(346436, 1_000, 3, "H2bw3Y2AjxJyK7smy1ZBB4LJ7MY3i9bPQM3YdAChAww2")]
+    #[test_case(454357, 10_000, 0, "4BLanrC5t7vzNXx62javKtjCmCkd8yZfZpVrjT4eUpNQ")]
+    #[test_case(454357, 10_000, 1, "FyvbdxpVchendERMnzH2KDceqydpXtJarrfFXoLQEXgQ")]
+    #[test_case(454357, 10_000, 2, "7KwK44Y7V3GzJLN8aGZtM8EEfAYmRvaiDyKYV6jg4MQn")]
+    #[test_case(454357, 10_000, 3, "E9XL5BLhCJ4Emyfs8jTUsQetfA8QZj78LcnN63dPp7jJ")]
+    fn test_long_leader_schedule_hashed(
+        epoch: Epoch,
+        len: u64,
+        stake_pow: u32,
+        expected_hash: &str,
+    ) {
+        fn hash_pubkeys(v: &[Pubkey]) -> String {
+            use sha2::{Digest, Sha256};
+
+            let hasher = v.iter().fold(Sha256::new(), |hasher, pk| {
+                hasher.chain_update(pk.to_bytes())
+            });
+            bs58::encode(hasher.finalize()).into_string()
+        }
+        let pubkeys: Vec<_> = (0..=u16::MAX).map(pubkey_from_u16).collect();
+        let stakes = pubkeys
+            .iter()
+            .enumerate()
+            .map(|(i, pk)| (pk, i.pow(stake_pow) as u64))
+            .collect();
+        let schedule = stake_weighted_slot_leaders(stakes, epoch, len, 1);
+        assert_eq!(hash_pubkeys(&schedule), expected_hash);
+    }
+
+    #[test]
+    #[should_panic]
+    fn test_zero_stake_panics() {
+        let _ = stake_weighted_slot_leaders(
+            vec![(&pubkey_from_u16(1), 0), (&pubkey_from_u16(2), 0)],
+            0,
+            5,
+            1,
+        );
+    }
 }