[test] Have real epoch_manager as default for tests

chain/client/src/sync/block.rs

@@ -369,7 +369,8 @@ mod test {
             BlockSync::new(network_adapter.clone().into(), block_fetch_horizon, false, true);
         let mut genesis_config = GenesisConfig::test();
         genesis_config.epoch_length = 100;
-        let mut env = TestEnv::builder(&genesis_config).clients_count(2).build();
+        let mut env =
+            TestEnv::builder(&genesis_config).clients_count(2).mock_epoch_managers().build();
         let mut blocks = vec![];
         for i in 1..5 * MAX_BLOCK_REQUESTS + 1 {
             let block = env.clients[0].produce_block(i as u64).unwrap().unwrap();
@@ -448,7 +449,8 @@ mod test {
             BlockSync::new(network_adapter.clone().into(), block_fetch_horizon, true, true);
         let mut genesis_config = GenesisConfig::test();
         genesis_config.epoch_length = 5;
-        let mut env = TestEnv::builder(&genesis_config).clients_count(2).build();
+        let mut env =
+            TestEnv::builder(&genesis_config).clients_count(2).mock_epoch_managers().build();
         let mut blocks = vec![];
         for i in 1..41 {
             let block = env.clients[0].produce_block(i).unwrap().unwrap();

chain/client/src/test_utils/test_env_builder.rs

@@ -199,29 +199,6 @@ impl TestEnvBuilder {
         }
     }
 
-    /// Specifies custom MockEpochManager for each client.  This allows us to
-    /// construct [`TestEnv`] with a custom implementation.
-    ///
-    /// The vector must have the same number of elements as they are clients
-    /// (one by default).  If that does not hold, [`Self::build`] method will
-    /// panic.
-    pub fn mock_epoch_managers(mut self, epoch_managers: Vec<Arc<MockEpochManager>>) -> Self {
-        assert_eq!(epoch_managers.len(), self.clients.len());
-        assert!(self.epoch_managers.is_none(), "Cannot override twice");
-        assert!(
-            self.num_shards.is_none(),
-            "Cannot set both num_shards and epoch_managers at the same time"
-        );
-        assert!(
-            self.shard_trackers.is_none(),
-            "Cannot override epoch_managers after shard_trackers"
-        );
-        assert!(self.runtimes.is_none(), "Cannot override epoch_managers after runtimes");
-        self.epoch_managers =
-            Some(epoch_managers.into_iter().map(|epoch_manager| epoch_manager.into()).collect());
-        self
-    }
-
     /// Specifies custom EpochManagerHandle for each client.  This allows us to
     /// construct [`TestEnv`] with a custom implementation.
     ///
@@ -245,12 +222,8 @@ impl TestEnvBuilder {
         self
     }
 
-    pub fn real_epoch_managers(self) -> Self {
-        self.real_epoch_managers_with_test_overrides(None)
-    }
-
     /// Constructs real EpochManager implementations for each instance.
-    pub fn real_epoch_managers_with_test_overrides(
+    pub fn epoch_managers_with_test_overrides(
         self,
         test_overrides: Option<AllEpochConfigTestOverrides>,
     ) -> Self {
@@ -273,10 +246,17 @@ impl TestEnvBuilder {
 
     /// Internal impl to make sure EpochManagers are initialized.
     fn ensure_epoch_managers(self) -> Self {
-        let mut ret = self.ensure_stores();
+        let ret = self.ensure_stores();
         if ret.epoch_managers.is_some() {
             return ret;
         }
+        ret.epoch_managers_with_test_overrides(None)
+    }
+
+    /// Constructs MockEpochManager implementations for each instance.
+    pub fn mock_epoch_managers(self) -> Self {
+        assert!(self.epoch_managers.is_none(), "Cannot override twice");
+        let mut ret = self.ensure_stores();
         let epoch_managers: Vec<EpochManagerKind> = (0..ret.clients.len())
             .map(|i| {
                 let vs = ValidatorSchedule::new_with_shards(ret.num_shards.unwrap_or(1))

chain/client/src/tests/doomslug.rs

@@ -19,7 +19,11 @@ use near_primitives::validator_signer::InMemoryValidatorSigner;
 fn test_processing_skips_on_forks() {
     init_test_logger();
 
-    let mut env = TestEnv::default_builder().clients_count(2).validator_seats(2).build();
+    let mut env = TestEnv::default_builder()
+        .clients_count(2)
+        .validator_seats(2)
+        .mock_epoch_managers()
+        .build();
     let b1 = env.clients[1].produce_block(1).unwrap().unwrap();
     let b2 = env.clients[0].produce_block(2).unwrap().unwrap();
     assert_eq!(b1.header().prev_hash(), b2.header().prev_hash());

chain/client/src/tests/process_blocks.rs

@@ -18,7 +18,7 @@ use std::sync::Arc;
 /// if the second block is not requested
 #[test]
 fn test_not_process_height_twice() {
-    let mut env = TestEnv::default_builder().build();
+    let mut env = TestEnv::default_builder().mock_epoch_managers().build();
     let block = env.clients[0].produce_block(1).unwrap().unwrap();
     // modify the block and resign it
     let mut duplicate_block = block.clone();
@@ -53,7 +53,7 @@ fn test_not_process_height_twice() {
 /// Test that if a block contains chunks with invalid shard_ids, the client will return error.
 #[test]
 fn test_bad_shard_id() {
-    let mut env = TestEnv::default_builder().num_shards(4).build();
+    let mut env = TestEnv::default_builder().num_shards(4).mock_epoch_managers().build();
     let prev_block = env.clients[0].produce_block(1).unwrap().unwrap();
     env.process_block(0, prev_block, Provenance::PRODUCED);
     let mut block = env.clients[0].produce_block(2).unwrap().unwrap(); // modify the block and resign it
@@ -98,7 +98,7 @@ fn test_bad_shard_id() {
 /// Test that if a block's content (vrf_value) is corrupted, the invalid block will not affect the node's block processing
 #[test]
 fn test_bad_block_content_vrf() {
-    let mut env = TestEnv::default_builder().num_shards(4).build();
+    let mut env = TestEnv::default_builder().num_shards(4).mock_epoch_managers().build();
     let prev_block = env.clients[0].produce_block(1).unwrap().unwrap();
     env.process_block(0, prev_block, Provenance::PRODUCED);
     let block = env.clients[0].produce_block(2).unwrap().unwrap();
@@ -122,7 +122,7 @@ fn test_bad_block_content_vrf() {
 /// Test that if a block's signature is corrupted, the invalid block will not affect the node's block processing
 #[test]
 fn test_bad_block_signature() {
-    let mut env = TestEnv::default_builder().num_shards(4).build();
+    let mut env = TestEnv::default_builder().num_shards(4).mock_epoch_managers().build();
     let prev_block = env.clients[0].produce_block(1).unwrap().unwrap();
     env.process_block(0, prev_block, Provenance::PRODUCED);
     let block = env.clients[0].produce_block(2).unwrap().unwrap();

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

@@ -28,10 +28,7 @@ fn benchmark_large_chunk_production_time() {
     let tx_size = 3 * mb;
 
     let genesis = Genesis::test(vec!["test0".parse().unwrap(), "test1".parse().unwrap()], 1);
-    let mut env = TestEnv::builder(&genesis.config)
-        .real_epoch_managers()
-        .nightshade_runtimes(&genesis)
-        .build();
+    let mut env = TestEnv::builder(&genesis.config).nightshade_runtimes(&genesis).build();
 
     let account_id = env.get_client_id(0).clone();
     let signer =

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

@@ -37,10 +37,7 @@ fn change_shard_id_to_invalid() {
     let epoch_length = 5000000;
     let mut genesis = Genesis::test(vec!["test0".parse().unwrap(), "test1".parse().unwrap()], 1);
     genesis.config.epoch_length = epoch_length;
-    let mut env = TestEnv::builder(&genesis.config)
-        .real_epoch_managers()
-        .nightshade_runtimes(&genesis)
-        .build();
+    let mut env = TestEnv::builder(&genesis.config).nightshade_runtimes(&genesis).build();
 
     let mut last_block = env.clients[0].chain.get_block_by_height(0).unwrap();
 
@@ -179,10 +176,7 @@ fn check_process_flipped_block_fails_on_bit(
     let epoch_length = 5000000;
     let mut genesis = Genesis::test(vec!["test0".parse().unwrap(), "test1".parse().unwrap()], 1);
     genesis.config.epoch_length = epoch_length;
-    let mut env = TestEnv::builder(&genesis.config)
-        .real_epoch_managers()
-        .nightshade_runtimes(&genesis)
-        .build();
+    let mut env = TestEnv::builder(&genesis.config).nightshade_runtimes(&genesis).build();
 
     let mut last_block = env.clients[0].chain.get_block_by_height(0).unwrap();
 

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

@@ -31,7 +31,7 @@ use nearcore::test_utils::TestEnvNightshadeSetupExt;
 /// TODO (#2445): Enable challenges when they are working correctly.
 #[test]
 fn test_block_with_challenges() {
-    let mut env = TestEnv::default_builder().build();
+    let mut env = TestEnv::default_builder().mock_epoch_managers().build();
     let genesis = env.clients[0].chain.get_block_by_height(0).unwrap();
 
     let mut block = env.clients[0].produce_block(1).unwrap().unwrap();
@@ -60,10 +60,7 @@ fn test_block_with_challenges() {
 #[test]
 fn test_invalid_chunk_state() {
     let genesis = Genesis::test(vec!["test0".parse().unwrap()], 1);
-    let mut env = TestEnv::builder(&genesis.config)
-        .real_epoch_managers()
-        .nightshade_runtimes(&genesis)
-        .build();
+    let mut env = TestEnv::builder(&genesis.config).nightshade_runtimes(&genesis).build();
     env.produce_block(0, 1);
     let block_hash = env.clients[0].chain.get_block_hash_by_height(1).unwrap();
 
@@ -86,7 +83,7 @@ fn test_invalid_chunk_state() {
 
 #[test]
 fn test_verify_block_double_sign_challenge() {
-    let mut env = TestEnv::default_builder().clients_count(2).build();
+    let mut env = TestEnv::default_builder().clients_count(2).mock_epoch_managers().build();
     env.produce_block(0, 1);
     let genesis = env.clients[0].chain.get_block_by_height(0).unwrap();
     let b1 = env.clients[0].produce_block(2).unwrap().unwrap();
@@ -195,7 +192,7 @@ fn create_invalid_proofs_chunk(client: &mut Client) -> (ProduceChunkResult, Bloc
 
 #[test]
 fn test_verify_chunk_invalid_proofs_challenge() {
-    let mut env = TestEnv::default_builder().build();
+    let mut env = TestEnv::default_builder().mock_epoch_managers().build();
     env.produce_block(0, 1);
     let (ProduceChunkResult { chunk, .. }, block) =
         create_invalid_proofs_chunk(&mut env.clients[0]);
@@ -208,7 +205,7 @@ fn test_verify_chunk_invalid_proofs_challenge() {
 
 #[test]
 fn test_verify_chunk_invalid_proofs_challenge_decoded_chunk() {
-    let mut env = TestEnv::default_builder().build();
+    let mut env = TestEnv::default_builder().mock_epoch_managers().build();
     env.produce_block(0, 1);
     let (ProduceChunkResult { chunk: encoded_chunk, .. }, block) =
         create_invalid_proofs_chunk(&mut env.clients[0]);
@@ -223,7 +220,7 @@ fn test_verify_chunk_invalid_proofs_challenge_decoded_chunk() {
 
 #[test]
 fn test_verify_chunk_proofs_malicious_challenge_no_changes() {
-    let mut env = TestEnv::default_builder().build();
+    let mut env = TestEnv::default_builder().mock_epoch_managers().build();
     env.produce_block(0, 1);
     // Valid chunk
     let (ProduceChunkResult { chunk, .. }, block) = create_chunk(&mut env.clients[0], None, None);
@@ -236,7 +233,7 @@ fn test_verify_chunk_proofs_malicious_challenge_no_changes() {
 
 #[test]
 fn test_verify_chunk_proofs_malicious_challenge_valid_order_transactions() {
-    let mut env = TestEnv::default_builder().build();
+    let mut env = TestEnv::default_builder().mock_epoch_managers().build();
     env.produce_block(0, 1);
 
     let genesis_hash = *env.clients[0].chain.genesis().hash();
@@ -272,7 +269,7 @@ fn test_verify_chunk_proofs_malicious_challenge_valid_order_transactions() {
 
 #[test]
 fn test_verify_chunk_proofs_challenge_transaction_order() {
-    let mut env = TestEnv::default_builder().build();
+    let mut env = TestEnv::default_builder().mock_epoch_managers().build();
     env.produce_block(0, 1);
 
     let genesis_hash = *env.clients[0].chain.genesis().hash();
@@ -334,10 +331,7 @@ fn challenge(
 fn test_verify_chunk_invalid_state_challenge() {
     let mut genesis = Genesis::test(vec!["test0".parse().unwrap(), "test1".parse().unwrap()], 1);
     genesis.config.min_gas_price = 0;
-    let mut env = TestEnv::builder(&genesis.config)
-        .real_epoch_managers()
-        .nightshade_runtimes(&genesis)
-        .build();
+    let mut env = TestEnv::builder(&genesis.config).nightshade_runtimes(&genesis).build();
     let signer = InMemorySigner::from_seed("test0".parse().unwrap(), KeyType::ED25519, "test0");
     let validator_signer = create_test_signer("test0");
     let genesis_hash = *env.clients[0].chain.genesis().hash();

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

@@ -121,10 +121,7 @@ fn test_storage_after_commit_of_cold_update() {
     let mut genesis = Genesis::test(vec![test0(), test1()], 1);
     genesis.config.epoch_length = epoch_length;
     genesis.config.min_gas_price = 0;
-    let mut env = TestEnv::builder(&genesis.config)
-        .real_epoch_managers()
-        .nightshade_runtimes(&genesis)
-        .build();
+    let mut env = TestEnv::builder(&genesis.config).nightshade_runtimes(&genesis).build();
 
     let (storage, ..) = create_test_node_storage_with_cold(DB_VERSION, DbKind::Hot);
     let cold_db = storage.cold_db().unwrap();
@@ -233,7 +230,6 @@ fn test_cold_db_head_update() {
     let cold_db = storage.cold_db().unwrap();
     let mut env = TestEnv::builder(&genesis.config)
         .stores(vec![hot_store.clone()])
-        .real_epoch_managers()
         .nightshade_runtimes(&genesis)
         .build();
 
@@ -267,10 +263,7 @@ fn test_cold_db_copy_with_height_skips() {
     let mut genesis = Genesis::test(vec![test0(), test1()], 1);
     genesis.config.epoch_length = epoch_length;
     genesis.config.min_gas_price = 0;
-    let mut env = TestEnv::builder(&genesis.config)
-        .real_epoch_managers()
-        .nightshade_runtimes(&genesis)
-        .build();
+    let mut env = TestEnv::builder(&genesis.config).nightshade_runtimes(&genesis).build();
 
     let (storage, ..) = create_test_node_storage_with_cold(DB_VERSION, DbKind::Hot);
     let cold_db = storage.cold_db().unwrap();
@@ -364,10 +357,7 @@ fn test_initial_copy_to_cold(batch_size: usize) {
 
     let mut genesis = Genesis::test(vec![test0(), test1()], 1);
     genesis.config.epoch_length = epoch_length;
-    let mut env = TestEnv::builder(&genesis.config)
-        .real_epoch_managers()
-        .nightshade_runtimes(&genesis)
-        .build();
+    let mut env = TestEnv::builder(&genesis.config).nightshade_runtimes(&genesis).build();
 
     let (storage, ..) = create_test_node_storage_with_cold(DB_VERSION, DbKind::Archive);
 
@@ -449,7 +439,6 @@ fn test_cold_loop_on_gc_boundary() {
         .archive(true)
         .save_trie_changes(true)
         .stores(vec![hot_store.clone()])
-        .real_epoch_managers()
         .nightshade_runtimes(&genesis)
         .build();
 

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

@@ -91,10 +91,7 @@ fn test_continuous_epoch_sync_info_population() {
 
     let mut genesis = Genesis::test(vec!["test0".parse().unwrap(), "test1".parse().unwrap()], 1);
     genesis.config.epoch_length = epoch_length;
-    let mut env = TestEnv::builder(&genesis.config)
-        .real_epoch_managers()
-        .nightshade_runtimes(&genesis)
-        .build();
+    let mut env = TestEnv::builder(&genesis.config).nightshade_runtimes(&genesis).build();
 
     let mut last_hash = *env.clients[0].chain.genesis().hash();
     let mut last_epoch_id = EpochId::default();
@@ -246,10 +243,7 @@ fn test_epoch_sync_data_hash_from_epoch_sync_info() {
 
     let mut genesis = Genesis::test(vec!["test0".parse().unwrap(), "test1".parse().unwrap()], 1);
     genesis.config.epoch_length = epoch_length;
-    let mut env = TestEnv::builder(&genesis.config)
-        .real_epoch_managers()
-        .nightshade_runtimes(&genesis)
-        .build();
+    let mut env = TestEnv::builder(&genesis.config).nightshade_runtimes(&genesis).build();
 
     let mut last_hash = *env.clients[0].chain.genesis().hash();
     let mut last_epoch_id = EpochId::default();
@@ -331,7 +325,6 @@ fn test_node_after_simulated_sync() {
         .clients_count(num_clients)
         .real_stores()
         .use_state_snapshots()
-        .real_epoch_managers()
         .nightshade_runtimes(&genesis)
         .build();