fix(upgradability): properly count stake

near · Jul 20, 2020 · c16b5f2 · c16b5f2
1 parent eb4963b
commit c16b5f2
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 2 deletions.
diff --git a/chain/epoch_manager/src/lib.rs b/chain/epoch_manager/src/lib.rs
@@ -228,7 +228,9 @@ impl EpochManager {
         let total_block_producer_stake: u128 = epoch_info
             .block_producers_settlement
             .iter()
-            .map(|id| epoch_info.validators[*id as usize].stake)
+            .collect::<HashSet<_>>()
+            .iter()
+            .map(|&id| epoch_info.validators[*id as usize].stake)
             .sum();
 
         let next_version = if let Some((&version, stake)) =
@@ -2934,4 +2936,37 @@ mod tests {
             PROTOCOL_VERSION
         );
     }
+
+    #[test]
+    fn test_protocol_version_switch_with_many_seats() {
+        let store = create_test_store();
+        let mut config = epoch_config(10, 1, 4, 0, 90, 60, 0);
+        config.num_block_producer_seats_per_shard = vec![10];
+        let amount_staked = 1_000_000;
+        let validators = vec![stake("test1", amount_staked), stake("test2", amount_staked / 5)];
+        let mut epoch_manager = EpochManager::new(
+            store.clone(),
+            config.clone(),
+            0,
+            default_reward_calculator(),
+            validators.clone(),
+        )
+        .unwrap();
+        let h = hash_range(50);
+        record_block(&mut epoch_manager, CryptoHash::default(), h[0], 0, vec![]);
+        let mut block_info1 = block_info(1, 1, h[0], h[0], h[0], vec![], DEFAULT_TOTAL_SUPPLY);
+        block_info1.latest_protocol_version = 0;
+        epoch_manager.record_block_info(&h[1], block_info1, [0; 32]).unwrap();
+        for i in 2..32 {
+            record_block(&mut epoch_manager, h[i - 1], h[i], i as u64, vec![]);
+        }
+        assert_eq!(
+            epoch_manager.get_epoch_info(&EpochId(h[10])).unwrap().protocol_version,
+            PROTOCOL_VERSION
+        );
+        assert_eq!(
+            epoch_manager.get_epoch_info(&EpochId(h[20])).unwrap().protocol_version,
+            PROTOCOL_VERSION
+        );
+    }
 }
diff --git a/pytest/tests/sanity/upgradable.py b/pytest/tests/sanity/upgradable.py
@@ -35,7 +35,7 @@ def main():
         "%snear-%s" % (near_root, stable_branch),
         "--home=%s" % node_root, "testnet", "--v", "4", "--prefix", "test"
     ])
-    genesis_config_changes = [("epoch_length", 20), ("block_producer_kickout_threshold", 80), ("chunk_producer_kickout_threshold", 80)]
+    genesis_config_changes = [("epoch_length", 20), ("num_block_producer_seats", 10), ("num_block_producer_seats_per_shard", [10]), ("block_producer_kickout_threshold", 80), ("chunk_producer_kickout_threshold", 80)]
     node_dirs = [os.path.join(node_root, 'test%d' % i) for i in range(4)]
     for i, node_dir in enumerate(node_dirs):
         cluster.apply_genesis_changes(node_dir, genesis_config_changes)