refactor(resharding) - Make some tests compatible with arbitrary Sha…

…rdIds part 2 (#12475)

Making more tests compatible with arbitrary ShardIds. There are still
some left but making progress. The ultimate goal is to make the default
shard layout to contain random shard ids.

chain/chain/src/runtime/tests.rs

@@ -739,8 +739,7 @@ fn test_state_sync() {
     env.step_default(vec![]);
     let block_hash = hash(&[env.head.height as u8]);
 
-    let shard_layout =
-        env.epoch_manager.get_shard_layout_from_prev_block(&env.head.prev_block_hash).unwrap();
+    let shard_layout = env.epoch_manager.get_shard_layout(&env.head.epoch_id).unwrap();
     let shard_id = shard_layout.shard_ids().next().unwrap();
 
     let state_part = env

chain/client/src/test_utils/client.rs

@@ -24,6 +24,7 @@ use near_primitives::transaction::SignedTransaction;
 use near_primitives::types::{BlockHeight, EpochId, ShardId};
 use near_primitives::utils::MaybeValidated;
 use near_primitives::version::PROTOCOL_VERSION;
+use near_store::ShardUId;
 use num_rational::Ratio;
 use reed_solomon_erasure::galois_8::ReedSolomon;
 
@@ -158,7 +159,7 @@ fn create_chunk_on_height_for_shard(
 }
 
 pub fn create_chunk_on_height(client: &mut Client, next_height: BlockHeight) -> ProduceChunkResult {
-    create_chunk_on_height_for_shard(client, next_height, ShardId::new(0))
+    create_chunk_on_height_for_shard(client, next_height, ShardUId::single_shard().shard_id())
 }
 
 pub fn create_chunk_with_transactions(

chain/client/src/test_utils/test_env.rs

@@ -197,6 +197,10 @@ impl TestEnv {
         let _ = cell.set(());
     }
 
+    pub fn contains_client(&self, account_id: &AccountId) -> bool {
+        self.account_indices.contains(account_id)
+    }
+
     pub fn client(&mut self, account_id: &AccountId) -> &mut Client {
         self.account_indices.lookup_mut(&mut self.clients, account_id)
     }
@@ -399,6 +403,10 @@ impl TestEnv {
                     let processing_done_tracker = ProcessingDoneTracker::new();
                     witness_processing_done_waiters.push(processing_done_tracker.make_waiter());
 
+                    if !self.contains_client(&account_id) {
+                        tracing::warn!(target: "test", "Client not found for account_id {}", account_id);
+                        continue;
+                    }
                     let client = self.client(&account_id);
                     let processing_result = client.process_chunk_state_witness(
                         state_witness.clone(),
@@ -434,6 +442,10 @@ impl TestEnv {
                     account_id,
                     endorsement,
                 )) => {
+                    if !self.contains_client(&account_id) {
+                        tracing::warn!(target: "test", "Client not found for account_id {}", account_id);
+                        return None;
+                    }
                     let processing_result = self
                         .client(&account_id)
                         .chunk_endorsement_tracker
@@ -877,6 +889,10 @@ impl AccountIndices {
         self.0[account_id]
     }
 
+    pub fn contains(&self, account_id: &AccountId) -> bool {
+        self.0.contains_key(account_id)
+    }
+
     pub fn lookup<'a, T>(&self, container: &'a [T], account_id: &AccountId) -> &'a T {
         &container[self.0[account_id]]
     }

core/primitives/src/shard_layout.rs

@@ -807,7 +807,7 @@ impl ShardUId {
     }
 
     /// Returns the only shard uid in the ShardLayout::single_shard layout.
-    /// It is not suitable for use with any other shard layouts.
+    /// It is not suitable for use with any other shard layout.
     pub fn single_shard() -> Self {
         ShardLayout::single_shard().shard_uids().next().unwrap()
     }

genesis-tools/genesis-populate/src/lib.rs

@@ -32,7 +32,7 @@ use near_vm_runner::logic::ProtocolVersion;
 use near_vm_runner::ContractCode;
 use nearcore::{NearConfig, NightshadeRuntime, NightshadeRuntimeExt};
 pub use node_runtime::bootstrap_congestion_info;
-use std::collections::BTreeMap;
+use std::collections::HashMap;
 use std::hash::{Hash, Hasher};
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
@@ -73,9 +73,9 @@ pub struct GenesisBuilder {
     store: Store,
     epoch_manager: Arc<EpochManagerHandle>,
     runtime: Arc<NightshadeRuntime>,
-    unflushed_records: BTreeMap<ShardId, Vec<StateRecord>>,
-    roots: BTreeMap<ShardId, StateRoot>,
-    state_updates: BTreeMap<ShardId, TrieUpdate>,
+    unflushed_records: HashMap<ShardId, Vec<StateRecord>>,
+    roots: HashMap<ShardId, StateRoot>,
+    state_updates: HashMap<ShardId, TrieUpdate>,
 
     // Things that can be set.
     additional_accounts_num: u64,
@@ -135,9 +135,15 @@ impl GenesisBuilder {
         // First, apply whatever is defined by the genesis config.
         let roots = get_genesis_state_roots(self.runtime.store())?
             .expect("genesis state roots not initialized.");
-        let genesis_shard_version = self.genesis.config.shard_layout.version();
-        self.roots =
-            roots.into_iter().enumerate().map(|(k, v)| (ShardId::new(k as u64), v)).collect();
+        let shard_layout = &self.genesis.config.shard_layout;
+        let genesis_shard_version = shard_layout.version();
+        self.roots = roots
+            .into_iter()
+            .enumerate()
+            .map(|(shard_index, state_root)| {
+                (shard_layout.get_shard_id(shard_index).unwrap(), state_root)
+            })
+            .collect();
         self.state_updates = self
             .roots
             .iter()

integration-tests/src/test_loop/tests/protocol_upgrade.rs

@@ -20,13 +20,6 @@ use crate::test_loop::builder::TestLoopBuilder;
 use crate::test_loop::env::TestLoopEnv;
 use crate::test_loop::utils::ONE_NEAR;
 
-// // Check that the shard ids and version are as expected. This is needed because
-// // the chunk drop condition uses those values.
-// fn assert_shard_layout(shard_layout: &ShardLayout) {
-//     assert_eq!(shard_layout.shard_ids().sorted().collect_vec(), vec![0, 1, 2, 3]);
-//     assert_eq!(shard_layout.version(), 1);
-// }
-
 /// Test upgrading the blockchain to another protocol version.
 /// Optionally make some chunks around epoch boundary missing.
 /// Uses a hardcoded shard layout, it doesn't change during the test.

integration-tests/src/tests/client/features/orphan_chunk_state_witness.rs

@@ -86,10 +86,11 @@ fn setup_orphan_witness_test() -> OrphanWitnessTestEnv {
     // The `excluded_validator` will receive a chunk witness for the chunk in `block2`, but it won't
     // have `block1`, so it will become an orphaned chunk state witness.
     let tip = env.clients[0].chain.head().unwrap();
+    let shard_layout = env.clients[0].epoch_manager.get_shard_layout(&tip.epoch_id).unwrap();
+    let shard_id = shard_layout.shard_ids().next().unwrap();
 
     let block1_producer = env.get_block_producer_at_offset(&tip, 1);
     let block2_producer = env.get_block_producer_at_offset(&tip, 2);
-    let shard_id = ShardId::new(0);
     let block2_chunk_producer = env.get_chunk_producer_at_offset(&tip, 2, shard_id);
 
     // The excluded validator shouldn't produce any blocks or chunks in the next two blocks.

integration-tests/src/tests/client/process_blocks.rs

@@ -57,9 +57,7 @@ use near_primitives::transaction::{
 };
 use near_primitives::trie_key::TrieKey;
 use near_primitives::types::validator_stake::ValidatorStake;
-use near_primitives::types::{
-    AccountId, BlockHeight, EpochId, NumBlocks, ProtocolVersion, ShardId,
-};
+use near_primitives::types::{AccountId, BlockHeight, EpochId, NumBlocks, ProtocolVersion};
 use near_primitives::version::{ProtocolFeature, PROTOCOL_VERSION};
 use near_primitives::views::{
     BlockHeaderView, FinalExecutionStatus, QueryRequest, QueryResponseKind,
@@ -1722,7 +1720,9 @@ fn ultra_slow_test_process_block_after_state_sync() {
         }
     };
     let sync_block = env.clients[0].chain.get_block(&sync_hash).unwrap();
-    let shard_id = ShardId::new(0);
+    let shard_layout =
+        env.clients[0].epoch_manager.get_shard_layout(sync_block.header().epoch_id()).unwrap();
+    let shard_id = shard_layout.shard_ids().next().unwrap();
 
     let header = env.clients[0].chain.compute_state_response_header(shard_id, sync_hash).unwrap();
     let state_root = header.chunk_prev_state_root();
@@ -2384,19 +2384,21 @@ fn test_validate_chunk_extra() {
         )
         .unwrap();
     let block = client.produce_block_on(next_height + 2, *block1.hash()).unwrap().unwrap();
+    let epoch_id = *block.header().epoch_id();
     client.process_block_test(block.into(), Provenance::PRODUCED).unwrap();
     env.propagate_chunk_state_witnesses_and_endorsements(false);
     let client = &mut env.clients[0];
+    let shard_layout = client.epoch_manager.get_shard_layout(&epoch_id).unwrap();
+    let shard_uid = shard_layout.shard_uids().next().unwrap();
+    let shard_id = shard_uid.shard_id();
     let chunks = client
         .chunk_inclusion_tracker
         .get_chunk_headers_ready_for_inclusion(block1.header().epoch_id(), &block1.hash());
-    let shard_id = ShardId::new(0);
     let (chunk_header, _) = client
         .chunk_inclusion_tracker
         .get_chunk_header_and_endorsements(chunks.get(&shard_id).unwrap())
         .unwrap();
-    let chunk_extra =
-        client.chain.get_chunk_extra(block1.hash(), &ShardUId::single_shard()).unwrap();
+    let chunk_extra = client.chain.get_chunk_extra(block1.hash(), &shard_uid).unwrap();
     assert!(validate_chunk_with_chunk_extra(
         &mut chain_store,
         client.epoch_manager.as_ref(),
@@ -2478,7 +2480,11 @@ fn slow_test_catchup_gas_price_change() {
     let sync_prev_block = &blocks[sync_block_idx - 1];
 
     assert!(env.clients[0].chain.check_sync_hash_validity(&sync_hash).unwrap());
-    let shard_id = ShardId::new(0);
+
+    let epoch_id = sync_prev_block.header().epoch_id();
+    let shard_layout = env.clients[0].epoch_manager.get_shard_layout(&epoch_id).unwrap();
+    let shard_id = shard_layout.shard_uids().next().unwrap().shard_id();
+
     let state_sync_header =
         env.clients[0].chain.get_state_response_header(shard_id, sync_hash).unwrap();
     let num_parts = state_sync_header.num_state_parts();
@@ -3708,13 +3714,14 @@ mod contract_precompilation_tests {
     const EPOCH_LENGTH: u64 = 25;
 
     fn state_sync_on_height(env: &TestEnv, height: BlockHeight) {
-        let shard_id = ShardId::new(0);
         let sync_block = env.clients[0].chain.get_block_by_height(height).unwrap();
         let sync_hash = *sync_block.hash();
+        let epoch_id = *env.clients[0].chain.get_block_header(&sync_hash).unwrap().epoch_id();
+        let shard_layout = env.clients[0].epoch_manager.get_shard_layout(&epoch_id).unwrap();
+        let shard_id = shard_layout.shard_ids().next().unwrap();
         let state_sync_header =
             env.clients[0].chain.get_state_response_header(shard_id, sync_hash).unwrap();
         let state_root = state_sync_header.chunk_prev_state_root();
-        let epoch_id = *env.clients[0].chain.get_block_header(&sync_hash).unwrap().epoch_id();
         let sync_prev_header =
             env.clients[0].chain.get_previous_header(sync_block.header()).unwrap();
         let sync_prev_prev_hash = sync_prev_header.prev_hash();

tools/state-viewer/src/state_dump.rs

@@ -318,6 +318,7 @@ mod test {
     use near_client::ProcessTxResponse;
     use near_crypto::{InMemorySigner, KeyFile, KeyType, PublicKey, SecretKey};
     use near_epoch_manager::EpochManager;
+    use near_o11y::testonly::init_test_logger;
     use near_primitives::account::id::AccountId;
     use near_primitives::state_record::StateRecord;
     use near_primitives::transaction::{Action, DeployContractAction, SignedTransaction};
@@ -797,6 +798,8 @@ mod test {
 
     #[test]
     fn test_dump_state_with_delayed_receipt() {
+        init_test_logger();
+
         let epoch_length = 4;
         let mut genesis =
             Genesis::test(vec!["test0".parse().unwrap(), "test1".parse().unwrap()], 1);