test: improve test loop builders

core/chain-configs/src/test_genesis.rs

@@ -1,5 +1,4 @@
-use std::collections::{BTreeMap, HashMap, HashSet};
-use std::sync::Arc;
+use std::collections::{HashMap, HashSet};
 
 use near_crypto::PublicKey;
 use near_primitives::account::{AccessKey, Account, AccountContract};
@@ -153,6 +152,20 @@ impl TestEpochConfigBuilder {
         Default::default()
     }
 
+    pub fn from_genesis(genesis: &Genesis) -> Self {
+        let mut builder = Self::new();
+        builder.epoch_length = genesis.config.epoch_length;
+        builder.shard_layout = genesis.config.shard_layout.clone();
+        builder.num_block_producer_seats = genesis.config.num_block_producer_seats;
+        builder.num_chunk_producer_seats = genesis.config.num_chunk_producer_seats;
+        builder.num_chunk_validator_seats = genesis.config.num_chunk_validator_seats;
+        builder
+    }
+
+    pub fn build_store_from_genesis(genesis: &Genesis) -> EpochConfigStore {
+        Self::from_genesis(genesis).build_store_for_single_version(genesis.config.protocol_version)
+    }
+
     pub fn epoch_length(mut self, epoch_length: BlockHeightDelta) -> Self {
         self.epoch_length = epoch_length;
         self
@@ -244,6 +257,14 @@ impl TestEpochConfigBuilder {
         tracing::debug!("Epoch config: {:#?}", epoch_config);
         epoch_config
     }
+
+    pub fn build_store_for_single_version(
+        self,
+        protocol_version: ProtocolVersion,
+    ) -> EpochConfigStore {
+        let epoch_config = self.build();
+        EpochConfigStore::test_single_version(protocol_version, epoch_config)
+    }
 }
 
 impl Default for TestGenesisBuilder {
@@ -671,21 +692,16 @@ pub fn build_genesis_and_epoch_config_store<'a>(
         .genesis_time_from_clock(&FakeClock::default().clock())
         .protocol_version(protocol_version)
         .epoch_length(epoch_length)
-        .shard_layout(shard_layout.clone())
-        .validators_spec(validators_spec.clone())
+        .shard_layout(shard_layout)
+        .validators_spec(validators_spec)
         .add_user_accounts_simple(accounts, 1_000_000 * ONE_NEAR)
         .gas_limit_one_petagas();
-    let epoch_config_builder = TestEpochConfigBuilder::new()
-        .epoch_length(epoch_length)
-        .shard_layout(shard_layout)
-        .validators_spec(validators_spec);
     let genesis_builder = customize_genesis_builder(genesis_builder);
-    let epoch_config_builder = customize_epoch_config_builder(epoch_config_builder);
-
     let genesis = genesis_builder.build();
-    let epoch_config = epoch_config_builder.build();
-    let epoch_config_store =
-        EpochConfigStore::test(BTreeMap::from([(protocol_version, Arc::new(epoch_config))]));
+
+    let epoch_config_builder = TestEpochConfigBuilder::from_genesis(&genesis);
+    let epoch_config_builder = customize_epoch_config_builder(epoch_config_builder);
+    let epoch_config_store = epoch_config_builder.build_store_for_single_version(protocol_version);
 
     (genesis, epoch_config_store)
 }

core/primitives/src/epoch_manager.rs

@@ -631,6 +631,13 @@ impl EpochConfigStore {
         Self { store }
     }
 
+    pub fn test_single_version(
+        protocol_version: ProtocolVersion,
+        epoch_config: EpochConfig,
+    ) -> Self {
+        Self::test(BTreeMap::from([(protocol_version, Arc::new(epoch_config))]))
+    }
+
     /// Returns the EpochConfig for the given protocol version.
     /// This panics if no config is found for the given version, thus the initialization via `for_chain_id` should
     /// only be performed for chains with some configs stored in files.

test-loop-tests/src/builder.rs

@@ -1,3 +1,4 @@
+use near_chain_configs::test_genesis::TestGenesisBuilder;
 use std::collections::{HashMap, HashSet};
 use std::sync::{Arc, Mutex};
 use tempfile::TempDir;
@@ -306,6 +307,12 @@ impl TestLoopBuilder {
         }
     }
 
+    // Creates TestLoop-compatible genesis builder
+    pub(crate) fn new_genesis_builder() -> TestGenesisBuilder {
+        TestGenesisBuilder::new()
+            .genesis_time_from_clock(&near_async::time::FakeClock::default().clock())
+    }
+
     /// Get the clock for the test loop.
     pub(crate) fn clock(&self) -> Clock {
         self.test_loop.clock()

test-loop-tests/src/tests/global_contracts.rs

@@ -1,17 +1,11 @@
-use std::collections::BTreeMap;
-use std::sync::Arc;
-
 use assert_matches::assert_matches;
 use near_async::time::Duration;
-use near_chain_configs::test_genesis::{
-    TestEpochConfigBuilder, TestGenesisBuilder, ValidatorsSpec,
-};
+use near_chain_configs::test_genesis::{TestEpochConfigBuilder, ValidatorsSpec};
 use near_client::Client;
 use near_client::test_utils::test_loop::ClientQueries;
 use near_o11y::testonly::init_test_logger;
 use near_parameters::{ActionCosts, RuntimeConfigStore, RuntimeFeesConfig};
 use near_primitives::action::{GlobalContractDeployMode, GlobalContractIdentifier};
-use near_primitives::epoch_manager::EpochConfigStore;
 use near_primitives::errors::{
     ActionError, ActionErrorKind, FunctionCallError, MethodResolveError, TxExecutionError,
 };
@@ -181,24 +175,16 @@ impl GlobalContractsTestEnv {
         let validators_spec =
             ValidatorsSpec::desired_roles(&block_and_chunk_producers, &chunk_validators_only);
 
-        let genesis = TestGenesisBuilder::new()
-            .genesis_time_from_clock(&near_async::time::FakeClock::default().clock())
-            .validators_spec(validators_spec.clone())
-            .shard_layout(shard_layout.clone())
+        let genesis = TestLoopBuilder::new_genesis_builder()
+            .validators_spec(validators_spec)
+            .shard_layout(shard_layout)
             .add_user_accounts_simple(
                 &[account_shard_0.clone(), account_shard_1.clone(), deploy_account.clone()],
                 initial_balance,
             )
             .gas_prices(GAS_PRICE, GAS_PRICE)
             .build();
-        let epoch_config = TestEpochConfigBuilder::new()
-            .shard_layout(shard_layout)
-            .validators_spec(validators_spec)
-            .build();
-        let epoch_config_store = EpochConfigStore::test(BTreeMap::from([(
-            genesis.config.protocol_version,
-            Arc::new(epoch_config),
-        )]));
+        let epoch_config_store = TestEpochConfigBuilder::build_store_from_genesis(&genesis);
 
         let clients = block_and_chunk_producers
             .iter()