impl

Author: pugachAG

core/chain-configs/src/test_genesis.rs

@@ -392,6 +392,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 +520,14 @@ impl TestGenesisBuilder {
             contents: GenesisContents::Records { records: GenesisRecords(records) },
         }
     }
+
+    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 {

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

@@ -181,24 +181,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())
+        let (genesis, epoch_config_store) = TestLoopBuilder::new_genesis_builder()
             .validators_spec(validators_spec.clone())
             .shard_layout(shard_layout.clone())
             .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()