impl

Author: pugachAG

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};
@@ -392,6 +391,14 @@ impl TestGenesisBuilder {
         self
     }
 
+    /// Creates epoch config builder compatible with the current genesis builder state
+    pub fn derive_epoch_config_builder(&self) -> TestEpochConfigBuilder {
+        TestEpochConfigBuilder::new()
+            .epoch_length(self.epoch_length)
+            .shard_layout(self.shard_layout.clone())
+            .validators_spec(self.validators_spec.clone())
+    }
+
     pub fn build(self) -> Genesis {
         if self
             .user_accounts
@@ -512,6 +519,16 @@ impl TestGenesisBuilder {
             contents: GenesisContents::Records { records: GenesisRecords(records) },
         }
     }
+
+    // Builds genesis along with epoch config store containing single epoch config
+    // for the genesis protocol version.
+    pub fn build_with_simple_epoch_config_store(self) -> (Genesis, EpochConfigStore) {
+        let epoch_config = self.derive_epoch_config_builder().build();
+        let genesis = self.build();
+        let epoch_config_store =
+            EpochConfigStore::test_single_version(genesis.config.protocol_version, epoch_config);
+        (genesis, epoch_config_store)
+    }
 }
 
 impl ValidatorsSpec {
@@ -671,21 +688,17 @@ 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 epoch_config_builder = genesis_builder.derive_epoch_config_builder();
     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_store = EpochConfigStore::test_single_version(protocol_version, epoch_config);
 
     (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::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,15 @@ 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, epoch_config_store) = 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),
-        )]));
+            .build_with_simple_epoch_config_store();
 
         let clients = block_and_chunk_producers
             .iter()