Fix: handle case where total stake is very small

chain/epoch-manager/src/validator_selection.rs

@@ -838,10 +838,10 @@ mod tests {
         // Given `epoch_info` and `proposals` above, the sample at a given height is deterministic.
         let height = 42;
         let expected_assignments = vec![
-            vec![(1, 300), (0, 300), (2, 300), (3, 60)],
-            vec![(0, 600), (2, 200), (1, 200)],
-            vec![(3, 200), (2, 300), (1, 100), (0, 400)],
-            vec![(2, 200), (4, 140), (1, 400), (0, 200)],
+            vec![(4, 56), (1, 168), (2, 300), (3, 84), (0, 364)],
+            vec![(3, 70), (1, 300), (4, 42), (2, 266), (0, 308)],
+            vec![(4, 42), (1, 238), (3, 42), (0, 450), (2, 196)],
+            vec![(2, 238), (1, 294), (3, 64), (0, 378)],
         ];
         assert_eq!(epoch_info.sample_chunk_validators(height), expected_assignments);
     }

core/primitives/src/validator_mandates/compute_price.rs

@@ -1,4 +1,8 @@
-use {super::ValidatorMandatesConfig, near_primitives_core::types::Balance, std::cmp::Ordering};
+use {
+    super::ValidatorMandatesConfig,
+    near_primitives_core::types::Balance,
+    std::cmp::{min, Ordering},
+};
 
 /// Given the stakes for the validators and the target number of mandates to have,
 /// this function computes the mandate price to use. It works by iterating a
@@ -18,7 +22,13 @@ where
 {
     let ValidatorMandatesConfig { target_mandates_per_shard, num_shards } = config;
     let total_stake = saturating_sum(stakes());
-    let target_mandates: u128 = num_shards.saturating_mul(target_mandates_per_shard) as u128;
+
+    // The target number of mandates cannot be larger than the total amount of stake.
+    // In production the total stake is _much_ higher than
+    // `num_shards * target_mandates_per_shard`, but in tests validators are given
+    // low staked numbers, so we need to have this condition in place.
+    let target_mandates: u128 =
+        min(num_shards.saturating_mul(target_mandates_per_shard) as u128, total_stake);
 
     let initial_price = total_stake / target_mandates;
 
@@ -96,6 +106,19 @@ mod tests {
 
     use super::*;
 
+    // Test case where the target number of mandates is larger than the total stake.
+    // This should never happen in production, but nearcore tests sometimes have
+    // low stake.
+    #[test]
+    fn test_small_total_stake() {
+        let stakes = [100_u128; 1];
+        let num_shards = 1;
+        let target_mandates_per_shard = 1000;
+        let config = ValidatorMandatesConfig::new(target_mandates_per_shard, num_shards);
+
+        assert_eq!(compute_mandate_price(config, || stakes.iter().copied()), 1);
+    }
+
     // Test cases where all stakes are equal.
     #[test]
     fn test_constant_dist() {