fix test_chunk_state_witness_bad_shard_id test

chain/client/src/view_client.rs

@@ -280,7 +280,7 @@ impl ViewClientActor {
             let cps: Vec<AccountId> = shard_ids
                 .iter()
                 .map(|&shard_id| {
-                    let cp = epoch_info.sample_chunk_producer(block_height, shard_id);
+                    let cp = epoch_info.sample_chunk_producer(block_height, shard_id).unwrap();
                     let cp = epoch_info.get_validator(cp).account_id().clone();
                     cp
                 })

chain/epoch-manager/src/lib.rs

@@ -1029,7 +1029,7 @@ impl EpochManager {
         shard_id: ShardId,
     ) -> Result<ValidatorStake, EpochError> {
         let epoch_info = self.get_epoch_info(epoch_id)?;
-        let validator_id = Self::chunk_producer_from_info(&epoch_info, height, shard_id);
+        let validator_id = Self::chunk_producer_from_info(&epoch_info, height, shard_id)?;
         Ok(epoch_info.get_validator(validator_id))
     }
 
@@ -1556,8 +1556,12 @@ impl EpochManager {
         epoch_info: &EpochInfo,
         height: BlockHeight,
         shard_id: ShardId,
-    ) -> ValidatorId {
-        epoch_info.sample_chunk_producer(height, shard_id)
+    ) -> Result<ValidatorId, EpochError> {
+        epoch_info.sample_chunk_producer(height, shard_id).ok_or_else(|| {
+            EpochError::ChunkProducerSelectionError(format!(
+                "Invalid shard {shard_id} for height {height}"
+            ))
+        })
     }
 
     /// Returns true, if given current block info, next block supposed to be in the next epoch.

chain/epoch-manager/src/tests/mod.rs

@@ -1131,7 +1131,7 @@ fn test_expected_chunks_prev_block_not_produced() {
         let prev_block_info = epoch_manager.get_block_info(&prev_block).unwrap();
         let prev_height = prev_block_info.height();
         let expected_chunk_producer =
-            EpochManager::chunk_producer_from_info(&epoch_info, prev_height + 1, 0);
+            EpochManager::chunk_producer_from_info(&epoch_info, prev_height + 1, 0).unwrap();
         // test1 does not produce blocks during first epoch
         if block_producer == 0 && epoch_id == initial_epoch_id {
             expected += 1;

chain/epoch-manager/src/types.rs

@@ -138,7 +138,8 @@ impl EpochInfoAggregator {
                 epoch_info,
                 prev_block_height + 1,
                 i as ShardId,
-            );
+            )
+            .unwrap();
             let tracker = self.shard_tracker.entry(i as ShardId).or_insert_with(HashMap::new);
             tracker
                 .entry(chunk_validator_id)

chain/epoch-manager/src/validator_selection.rs

@@ -606,7 +606,7 @@ mod tests {
                 let cp = epoch_info.sample_chunk_producer(h, shard_id);
                 // Don't read too much into this. The reason the ValidatorId always
                 // equals the ShardId is because the validators are assigned to shards in order.
-                assert_eq!(cp, shard_id)
+                assert_eq!(cp, Some(shard_id))
             }
         }
 

core/primitives/src/epoch_manager.rs

@@ -1081,33 +1081,37 @@ pub mod epoch_info {
             }
         }
 
-        pub fn sample_chunk_producer(&self, height: BlockHeight, shard_id: ShardId) -> ValidatorId {
+        pub fn sample_chunk_producer(
+            &self,
+            height: BlockHeight,
+            shard_id: ShardId,
+        ) -> Option<ValidatorId> {
             match &self {
                 Self::V1(v1) => {
                     let cp_settlement = &v1.chunk_producers_settlement;
-                    let shard_cps = &cp_settlement[shard_id as usize];
-                    shard_cps[(height as u64 % (shard_cps.len() as u64)) as usize]
+                    let shard_cps = cp_settlement.get(shard_id as usize)?;
+                    shard_cps.get((height as u64 % (shard_cps.len() as u64)) as usize).copied()
                 }
                 Self::V2(v2) => {
                     let cp_settlement = &v2.chunk_producers_settlement;
-                    let shard_cps = &cp_settlement[shard_id as usize];
-                    shard_cps[(height as u64 % (shard_cps.len() as u64)) as usize]
+                    let shard_cps = cp_settlement.get(shard_id as usize)?;
+                    shard_cps.get((height as u64 % (shard_cps.len() as u64)) as usize).copied()
                 }
                 Self::V3(v3) => {
                     let protocol_version = self.protocol_version();
                     let seed =
                         Self::chunk_produce_seed(protocol_version, &v3.rng_seed, height, shard_id);
                     let shard_id = shard_id as usize;
-                    let sample = v3.chunk_producers_sampler[shard_id].sample(seed);
-                    v3.chunk_producers_settlement[shard_id][sample]
+                    let sample = v3.chunk_producers_sampler.get(shard_id)?.sample(seed);
+                    v3.chunk_producers_settlement.get(shard_id)?.get(sample).copied()
                 }
                 Self::V4(v4) => {
                     let protocol_version = self.protocol_version();
                     let seed =
                         Self::chunk_produce_seed(protocol_version, &v4.rng_seed, height, shard_id);
                     let shard_id = shard_id as usize;
-                    let sample = v4.chunk_producers_sampler[shard_id].sample(seed);
-                    v4.chunk_producers_settlement[shard_id][sample]
+                    let sample = v4.chunk_producers_sampler.get(shard_id)?.sample(seed);
+                    v4.chunk_producers_settlement.get(shard_id)?.get(sample).copied()
                 }
             }
         }

core/primitives/src/errors.rs

@@ -865,6 +865,8 @@ pub enum EpochError {
     },
     /// Error selecting validators for a chunk.
     ChunkValidatorSelectionError(String),
+    /// Error selecting chunk producer for a shard.
+    ChunkProducerSelectionError(String),
 }
 
 impl std::error::Error for EpochError {}
@@ -892,6 +894,9 @@ impl Display for EpochError {
             EpochError::ChunkValidatorSelectionError(err) => {
                 write!(f, "Error selecting validators for a chunk: {}", err)
             }
+            EpochError::ChunkProducerSelectionError(err) => {
+                write!(f, "Error selecting chunk producer: {}", err)
+            }
         }
     }
 }
@@ -915,6 +920,9 @@ impl Debug for EpochError {
             EpochError::ChunkValidatorSelectionError(err) => {
                 write!(f, "ChunkValidatorSelectionError({})", err)
             }
+            EpochError::ChunkProducerSelectionError(err) => {
+                write!(f, "ChunkProducerSelectionError({})", err)
+            }
         }
     }
 }