refactor: disentangle test genesis and epoch config building (#12554)

closes #12548 

Any inputs before I start to update the tests?

core/primitives/src/shard_layout.rs

@@ -391,6 +391,14 @@ impl ShardLayout {
         })
     }
 
+    /// Test-only helper to create a simple multi-shard ShardLayout with the provided boundaries.
+    /// The shard ids are deterministic but arbitrary in order to test the non-contiguous ShardIds.
+    pub fn simple_v1(boundary_accounts: &[&str]) -> ShardLayout {
+        // TODO these test methods should go into a different namespace
+        let boundary_accounts = boundary_accounts.iter().map(|a| a.parse().unwrap()).collect();
+        Self::multi_shard_custom(boundary_accounts, 1)
+    }
+
     /// Return a V0 Shardlayout
     #[deprecated(note = "Use multi_shard() instead")]
     pub fn v0(num_shards: NumShards, version: ShardVersion) -> Self {

integration-tests/src/test_loop/builder.rs

@@ -1,5 +1,6 @@
 use std::collections::{HashMap, HashSet};
 use std::sync::{Arc, Mutex};
+use tempfile::TempDir;
 
 use near_async::futures::FutureSpawner;
 use near_async::messaging::{noop, IntoMultiSender, IntoSender, LateBoundSender};
@@ -40,7 +41,6 @@ use near_store::{ShardUId, Store, StoreConfig, TrieConfig};
 use near_vm_runner::logic::ProtocolVersion;
 use near_vm_runner::{ContractRuntimeCache, FilesystemContractRuntimeCache};
 use nearcore::state_sync::StateSyncDumper;
-use tempfile::TempDir;
 
 use super::env::{ClientToShardsManagerSender, TestData, TestLoopChunksStorage, TestLoopEnv};
 use super::utils::network::{chunk_endorsement_dropper, chunk_endorsement_dropper_by_hash};

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

@@ -11,7 +11,9 @@ use near_async::test_loop::sender::TestLoopSender;
 use near_async::test_loop::TestLoopV2;
 use near_async::time::Duration;
 use near_chain::ChainStoreAccess;
-use near_chain_configs::test_genesis::TestGenesisBuilder;
+use near_chain_configs::test_genesis::{
+    build_genesis_and_epoch_config_store, GenesisAndEpochConfigParams, ValidatorsSpec,
+};
 use near_client::client_actor::ClientActorInner;
 use near_client::Client;
 use near_crypto::Signer;
@@ -44,7 +46,7 @@ use testlib::bandwidth_scheduler::get_random_receipt_size_for_test;
 use crate::test_loop::builder::TestLoopBuilder;
 use crate::test_loop::env::{TestData, TestLoopEnv};
 use crate::test_loop::utils::transactions::{run_txs_parallel, TransactionRunner};
-use crate::test_loop::utils::{ONE_NEAR, TGAS};
+use crate::test_loop::utils::TGAS;
 
 /// 1 node, 3 shards
 /// Lots of transactions which generate congestion and buffered cross-shard receipts.
@@ -85,24 +87,22 @@ fn slow_test_bandwidth_scheduler_request_generation() {
     all_accounts.extend(workload_accounts.clone());
     all_accounts.push(node_account.clone());
 
-    let builder = TestLoopBuilder::new();
-    let mut genesis_builder = TestGenesisBuilder::new();
-    genesis_builder
-        .genesis_time_from_clock(&builder.clock())
-        .protocol_version_latest()
-        .genesis_height(10000)
-        .gas_limit_one_petagas()
-        .shard_layout(shard_layout)
-        .transaction_validity_period(1000)
-        .epoch_length(10000)
-        .validators_desired_roles(&[node_account.as_str()], &[]);
-    for account in all_accounts {
-        genesis_builder.add_user_account_simple(account.clone(), 100_000_0000 * ONE_NEAR);
-    }
-
-    let (genesis, epoch_config_store) = genesis_builder.build();
-
-    let TestLoopEnv { mut test_loop, datas: node_datas, tempdir } = builder
+    let epoch_length = 10000;
+    let validators_spec = ValidatorsSpec::desired_roles(&[node_account.as_str()], &[]);
+
+    let (genesis, epoch_config_store) = build_genesis_and_epoch_config_store(
+        GenesisAndEpochConfigParams {
+            epoch_length,
+            protocol_version: PROTOCOL_VERSION,
+            shard_layout,
+            validators_spec,
+            accounts: &all_accounts,
+        },
+        |genesis_builder| genesis_builder.genesis_height(10000).transaction_validity_period(1000),
+        |epoch_config_builder| epoch_config_builder,
+    );
+
+    let TestLoopEnv { mut test_loop, datas: node_datas, tempdir } = TestLoopBuilder::new()
         .genesis(genesis)
         .epoch_config_store(epoch_config_store)
         .clients(vec![node_account])

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

@@ -1,13 +1,16 @@
 use crate::test_loop::builder::TestLoopBuilder;
 use crate::test_loop::env::TestLoopEnv;
 use crate::test_loop::utils::validators::get_epoch_all_validators;
-use crate::test_loop::utils::ONE_NEAR;
 use itertools::Itertools;
 use near_async::test_loop::data::TestLoopData;
 use near_async::time::Duration;
-use near_chain_configs::test_genesis::TestGenesisBuilder;
+use near_chain_configs::test_genesis::{
+    build_genesis_and_epoch_config_store, GenesisAndEpochConfigParams, ValidatorsSpec,
+};
 use near_o11y::testonly::init_test_logger;
+use near_primitives::shard_layout::ShardLayout;
 use near_primitives::types::AccountId;
+use near_primitives::version::PROTOCOL_VERSION;
 
 const NUM_ACCOUNTS: usize = 8;
 const NUM_PRODUCER_ACCOUNTS: usize = 6;
@@ -34,7 +37,6 @@ impl TestCase {
 
 fn run_test_chunk_validator_kickout(accounts: Vec<AccountId>, test_case: TestCase) {
     init_test_logger();
-    let initial_balance = 10000 * ONE_NEAR;
     let epoch_length = 10;
     let clients = accounts.iter().cloned().collect_vec();
     let accounts_str = accounts.iter().map(|a| a.as_str()).collect_vec();
@@ -71,20 +73,26 @@ fn run_test_chunk_validator_kickout(accounts: Vec<AccountId>, test_case: TestCas
             None
         };
 
-    let mut genesis_builder = TestGenesisBuilder::new();
-    genesis_builder
-        .genesis_time_from_clock(&builder.clock())
-        .shard_layout_simple_v1(&["account2", "account4", "account6"])
-        .epoch_length(epoch_length)
-        // Select 6 block&chunk producers and 2 chunk validators.
-        .validators_desired_roles(block_and_chunk_producers, chunk_validators_only)
-        // Set up config to kick out only chunk validators for low performance.
-        .kickouts_for_chunk_validators_only()
-        .target_validator_mandates_per_shard(num_validator_mandates_per_shard);
-    for account in &accounts {
-        genesis_builder.add_user_account_simple(account.clone(), initial_balance);
-    }
-    let (genesis, epoch_config_store) = genesis_builder.build();
+    let shard_layout = ShardLayout::simple_v1(&["account2", "account4", "account6"]);
+    let validators_spec =
+        ValidatorsSpec::desired_roles(block_and_chunk_producers, chunk_validators_only);
+
+    let (genesis, epoch_config_store) = build_genesis_and_epoch_config_store(
+        GenesisAndEpochConfigParams {
+            epoch_length,
+            protocol_version: PROTOCOL_VERSION,
+            shard_layout,
+            validators_spec,
+            accounts: &accounts,
+        },
+        |genesis_builder| genesis_builder,
+        |epoch_config_builder| {
+            epoch_config_builder
+                // Set up config to kick out only chunk validators for low performance.
+                .kickouts_for_chunk_validators_only()
+                .target_validator_mandates_per_shard(num_validator_mandates_per_shard)
+        },
+    );
 
     let TestLoopEnv { mut test_loop, datas: node_datas, tempdir } =
         builder.genesis(genesis).epoch_config_store(epoch_config_store).clients(clients).build();

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

@@ -4,17 +4,21 @@ use itertools::Itertools;
 use near_async::test_loop::data::{TestLoopData, TestLoopDataHandle};
 use near_async::test_loop::TestLoopV2;
 use near_async::time::Duration;
-use near_chain_configs::test_genesis::TestGenesisBuilder;
+use near_chain_configs::test_genesis::{
+    build_genesis_and_epoch_config_store, GenesisAndEpochConfigParams, ValidatorsSpec,
+};
 use near_client::client_actor::ClientActorInner;
 use near_o11y::testonly::init_test_logger;
+use near_primitives::shard_layout::ShardLayout;
 use near_primitives::types::{AccountId, BlockHeight};
+use near_primitives::version::PROTOCOL_VERSION;
 
 use crate::test_loop::builder::TestLoopBuilder;
 use crate::test_loop::env::{TestData, TestLoopEnv};
 use crate::test_loop::utils::transactions::{
     call_contract, check_txs, deploy_contract, make_accounts,
 };
-use crate::test_loop::utils::{ONE_NEAR, TGAS};
+use crate::test_loop::utils::TGAS;
 
 const NUM_ACCOUNTS: usize = 100;
 const NUM_PRODUCERS: usize = 2;
@@ -60,7 +64,6 @@ fn test_congestion_control_simple() {
 }
 
 fn setup(accounts: &Vec<AccountId>) -> (TestLoopEnv, AccountId) {
-    let initial_balance = 10000 * ONE_NEAR;
     let clients = accounts.iter().take(NUM_CLIENTS).cloned().collect_vec();
 
     // split the clients into producers, validators, and rpc nodes
@@ -74,26 +77,29 @@ fn setup(accounts: &Vec<AccountId>) -> (TestLoopEnv, AccountId) {
     let validators = validators.iter().map(|account| account.as_str()).collect_vec();
     let [rpc_id] = rpcs else { panic!("Expected exactly one rpc node") };
 
-    let builder = TestLoopBuilder::new();
-    let mut genesis_builder = TestGenesisBuilder::new();
-    genesis_builder
-        .genesis_time_from_clock(&builder.clock())
-        .protocol_version_latest()
-        .genesis_height(10000)
-        .gas_prices_free()
-        .gas_limit_one_petagas()
-        .shard_layout_simple_v1(&["account3", "account5", "account7"])
-        .transaction_validity_period(1000)
-        .epoch_length(10)
-        .validators_desired_roles(&producers, &validators)
-        .shuffle_shard_assignment_for_chunk_producers(true);
-    for account in accounts {
-        genesis_builder.add_user_account_simple(account.clone(), initial_balance);
-    }
-    let (genesis, epoch_config_store) = genesis_builder.build();
+    let epoch_length = 10;
+    let shard_layout = ShardLayout::simple_v1(&["account3", "account5", "account7"]);
+    let validators_spec = ValidatorsSpec::desired_roles(&producers, &validators);
+
+    let (genesis, epoch_config_store) = build_genesis_and_epoch_config_store(
+        GenesisAndEpochConfigParams {
+            epoch_length,
+            protocol_version: PROTOCOL_VERSION,
+            shard_layout,
+            validators_spec,
+            accounts: &accounts,
+        },
+        |genesis_builder| genesis_builder.genesis_height(10000).transaction_validity_period(1000),
+        |epoch_config_builder| {
+            epoch_config_builder.shuffle_shard_assignment_for_chunk_producers(true)
+        },
+    );
 
-    let env =
-        builder.genesis(genesis).epoch_config_store(epoch_config_store).clients(clients).build();
+    let env = TestLoopBuilder::new()
+        .genesis(genesis)
+        .epoch_config_store(epoch_config_store)
+        .clients(clients)
+        .build();
     (env, rpc_id.clone())
 }
 

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

@@ -1,14 +1,16 @@
 use near_async::time::Duration;
 use near_chain::ChainStoreAccess;
-use near_chain_configs::test_genesis::TestGenesisBuilder;
+use near_chain_configs::test_genesis::{
+    build_genesis_and_epoch_config_store, GenesisAndEpochConfigParams, ValidatorsSpec,
+};
 use near_client::Client;
 use near_o11y::testonly::init_test_logger;
+use near_primitives::shard_layout::ShardLayout;
 use near_primitives::types::AccountId;
 use near_primitives::version::{ProtocolFeature, PROTOCOL_VERSION};
 
 use crate::test_loop::builder::TestLoopBuilder;
 use crate::test_loop::env::TestLoopEnv;
-use crate::test_loop::utils::ONE_NEAR;
 
 const NUM_SHARDS: usize = 4;
 
@@ -25,23 +27,25 @@ fn test_congestion_control_genesis_bootstrap() {
 
     let builder = TestLoopBuilder::new();
 
-    let initial_balance = 10000 * ONE_NEAR;
     let accounts = ["test0", "test1"];
     let clients: Vec<AccountId> = accounts.iter().map(|account| account.parse().unwrap()).collect();
 
-    let mut genesis_builder = TestGenesisBuilder::new();
-    genesis_builder
-        .genesis_time_from_clock(&builder.clock())
-        .protocol_version_latest()
-        .shard_layout_simple_v1(&["account3", "account5", "account7"])
-        .validators_desired_roles(&accounts[0..1], &accounts[1..2])
-        .minimum_validators_per_shard(1);
+    let epoch_length = 100;
+    let shard_layout = ShardLayout::simple_v1(&["account3", "account5", "account7"]);
+    let validators_spec = ValidatorsSpec::desired_roles(&accounts[0..1], &accounts[1..2]);
 
-    for i in 0..clients.len() {
-        genesis_builder.add_user_account_simple(clients[i].clone(), initial_balance);
-    }
+    let (genesis, epoch_config_store) = build_genesis_and_epoch_config_store(
+        GenesisAndEpochConfigParams {
+            epoch_length,
+            protocol_version: PROTOCOL_VERSION,
+            shard_layout,
+            validators_spec,
+            accounts: &clients,
+        },
+        |genesis_builder| genesis_builder,
+        |epoch_config_builder| epoch_config_builder.minimum_validators_per_shard(1),
+    );
 
-    let (genesis, epoch_config_store) = genesis_builder.build();
     let TestLoopEnv { mut test_loop, datas: node_datas, tempdir } = builder
         .genesis(genesis)
         .epoch_config_store(epoch_config_store)

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

@@ -1,8 +1,12 @@
 use itertools::Itertools;
 use near_async::time::Duration;
-use near_chain_configs::test_genesis::TestGenesisBuilder;
+use near_chain_configs::test_genesis::{
+    build_genesis_and_epoch_config_store, GenesisAndEpochConfigParams, ValidatorsSpec,
+};
 use near_o11y::testonly::init_test_logger;
+use near_primitives::shard_layout::ShardLayout;
 use near_primitives::types::AccountId;
+use near_primitives::version::PROTOCOL_VERSION;
 use near_vm_runner::ContractCode;
 
 use crate::test_loop::builder::TestLoopBuilder;
@@ -11,10 +15,10 @@ use crate::test_loop::utils::contract_distribution::{
     assert_all_chunk_endorsements_received, clear_compiled_contract_caches,
     run_until_caches_contain_contract,
 };
+use crate::test_loop::utils::get_head_height;
 use crate::test_loop::utils::transactions::{
     call_contract, check_txs, deploy_contract, make_accounts,
 };
-use crate::test_loop::utils::{get_head_height, ONE_NEAR};
 
 const EPOCH_LENGTH: u64 = 10;
 const GENESIS_HEIGHT: u64 = 1000;
@@ -73,7 +77,6 @@ fn test_contract_distribution_cross_shard() {
 fn setup(accounts: &Vec<AccountId>) -> (TestLoopEnv, AccountId) {
     let builder = TestLoopBuilder::new();
 
-    let initial_balance = 10000 * ONE_NEAR;
     // All block_and_chunk_producers will be both block and chunk validators.
     let block_and_chunk_producers =
         (0..NUM_BLOCK_AND_CHUNK_PRODUCERS).map(|idx| accounts[idx].as_str()).collect_vec();
@@ -85,23 +88,27 @@ fn setup(accounts: &Vec<AccountId>) -> (TestLoopEnv, AccountId) {
     let clients = accounts.iter().take(NUM_VALIDATORS + NUM_RPC).cloned().collect_vec();
     let rpc_id = accounts[NUM_VALIDATORS].clone();
 
-    let mut genesis_builder = TestGenesisBuilder::new();
-    genesis_builder
-        .genesis_time_from_clock(&builder.clock())
-        .protocol_version_latest()
-        .genesis_height(GENESIS_HEIGHT)
-        .gas_prices_free()
-        .gas_limit_one_petagas()
-        .shard_layout_simple_v1(&["account4"])
-        .transaction_validity_period(1000)
-        .epoch_length(EPOCH_LENGTH)
-        .validators_desired_roles(&block_and_chunk_producers, &chunk_validators_only)
-        .shuffle_shard_assignment_for_chunk_producers(true)
-        .minimum_validators_per_shard(2);
-    for account in accounts {
-        genesis_builder.add_user_account_simple(account.clone(), initial_balance);
-    }
-    let (genesis, epoch_config_store) = genesis_builder.build();
+    let shard_layout = ShardLayout::simple_v1(&["account4"]);
+    let validators_spec =
+        ValidatorsSpec::desired_roles(&block_and_chunk_producers, &chunk_validators_only);
+
+    let (genesis, epoch_config_store) = build_genesis_and_epoch_config_store(
+        GenesisAndEpochConfigParams {
+            epoch_length: EPOCH_LENGTH,
+            protocol_version: PROTOCOL_VERSION,
+            shard_layout,
+            validators_spec,
+            accounts: &accounts,
+        },
+        |genesis_builder| {
+            genesis_builder.genesis_height(GENESIS_HEIGHT).transaction_validity_period(1000)
+        },
+        |epoch_config_builder| {
+            epoch_config_builder
+                .shuffle_shard_assignment_for_chunk_producers(true)
+                .minimum_validators_per_shard(2)
+        },
+    );
 
     let env =
         builder.genesis(genesis).epoch_config_store(epoch_config_store).clients(clients).build();