diff --git a/chain/chain/src/test_utils.rs b/chain/chain/src/test_utils.rs
index 6da9b859174..4778a8abcdb 100644
--- a/chain/chain/src/test_utils.rs
+++ b/chain/chain/src/test_utils.rs
@@ -16,7 +16,7 @@ use near_async::time::Clock;
 use near_chain_configs::Genesis;
 use near_chain_primitives::Error;
 use near_epoch_manager::shard_tracker::ShardTracker;
-use near_epoch_manager::EpochManager;
+use near_epoch_manager::{EpochManager, EpochManagerHandle};
 use near_primitives::block::Block;
 use near_primitives::hash::CryptoHash;
 use near_primitives::test_utils::create_test_signer;
@@ -124,7 +124,7 @@ pub fn process_block_sync(
 // TODO(#8190) Improve this testing API.
 pub fn setup(
     clock: Clock,
-) -> (Chain, Arc<MockEpochManager>, Arc<KeyValueRuntime>, Arc<InMemoryValidatorSigner>) {
+) -> (Chain, Arc<EpochManagerHandle>, Arc<NightshadeRuntime>, Arc<InMemoryValidatorSigner>) {
     setup_with_tx_validity_period(clock, 100, 1000)
 }
 
@@ -132,28 +132,34 @@ pub fn setup_with_tx_validity_period(
     clock: Clock,
     tx_validity_period: NumBlocks,
     epoch_length: u64,
-) -> (Chain, Arc<MockEpochManager>, Arc<KeyValueRuntime>, Arc<InMemoryValidatorSigner>) {
+) -> (Chain, Arc<EpochManagerHandle>, Arc<NightshadeRuntime>, Arc<InMemoryValidatorSigner>) {
     let store = create_test_store();
-    let epoch_manager = MockEpochManager::new(store.clone(), epoch_length);
+    let mut genesis = Genesis::test_sharded(
+        clock.clone(),
+        vec!["test".parse::<AccountId>().unwrap()],
+        1,
+        vec![1; 1],
+    );
+    genesis.config.epoch_length = epoch_length;
+    genesis.config.transaction_validity_period = tx_validity_period;
+    genesis.config.gas_limit = 1_000_000;
+    genesis.config.min_gas_price = 100;
+    genesis.config.max_gas_price = 1_000_000_000;
+    genesis.config.total_supply = 1_000_000_000;
+    genesis.config.gas_price_adjustment_rate = Ratio::from_integer(0);
+    genesis.config.protocol_version = PROTOCOL_VERSION;
+    let tempdir = tempfile::tempdir().unwrap();
+    initialize_genesis_state(store.clone(), &genesis, Some(tempdir.path()));
+    let epoch_manager = EpochManager::new_arc_handle(store.clone(), &genesis.config);
     let shard_tracker = ShardTracker::new_empty(epoch_manager.clone());
-    let runtime = KeyValueRuntime::new(store, epoch_manager.as_ref());
+    let runtime =
+        NightshadeRuntime::test(tempdir.path(), store, &genesis.config, epoch_manager.clone());
     let chain = Chain::new(
-        clock.clone(),
+        clock,
         epoch_manager.clone(),
         shard_tracker,
         runtime.clone(),
-        &ChainGenesis {
-            time: clock.now_utc(),
-            height: 0,
-            gas_limit: 1_000_000,
-            min_gas_price: 100,
-            max_gas_price: 1_000_000_000,
-            total_supply: 1_000_000_000,
-            gas_price_adjustment_rate: Ratio::from_integer(0),
-            transaction_validity_period: tx_validity_period,
-            epoch_length,
-            protocol_version: PROTOCOL_VERSION,
-        },
+        &ChainGenesis::new(&genesis.config),
         DoomslugThresholdMode::NoApprovals,
         ChainConfig::test(),
         None,
diff --git a/chain/chain/src/tests/simple_chain.rs b/chain/chain/src/tests/simple_chain.rs
index 099c8baaef9..5fce22e8d41 100644
--- a/chain/chain/src/tests/simple_chain.rs
+++ b/chain/chain/src/tests/simple_chain.rs
@@ -33,9 +33,9 @@ fn build_chain() {
     //     cargo insta test --accept -p near-chain --features nightly -- tests::simple_chain::build_chain
     let hash = chain.head().unwrap().last_block_hash;
     if cfg!(feature = "nightly") {
-        insta::assert_display_snapshot!(hash, @"9JRdRkSfVRwYBVFNzfKwuhuUabMkQN14b1Ge7z1FRnMF");
+        insta::assert_display_snapshot!(hash, @"Dnn7UkUiRfMu13jheAbcoJvQ66gKBgAoynRdEkMTkX58");
     } else {
-        insta::assert_display_snapshot!(hash, @"387773cK74c4cHxNWnrTgoZ2dkEhVjtFjJNpSxeWohPx");
+        insta::assert_display_snapshot!(hash, @"EsUNazp4zR2XgiwZSuQnX9dsaFk1VDhdRwGYt1YHpu5b");
     }
 
     for i in 1..5 {
@@ -51,9 +51,9 @@ fn build_chain() {
 
     let hash = chain.head().unwrap().last_block_hash;
     if cfg!(feature = "nightly") {
-        insta::assert_display_snapshot!(hash, @"HkPEDHWJHuW8Yzw9FonrJarVAHi51mQt7A4DQUtBagko");
+        insta::assert_display_snapshot!(hash, @"8xU6fbbdGeYmQRaAbsszMHDM88KSAaniHykowTUDYXpm");
     } else {
-        insta::assert_display_snapshot!(hash, @"6QpFm2bFYyfBjrk1KACMLr4dsezfZNYWGFnsbuD7gssH");
+        insta::assert_display_snapshot!(hash, @"CJ5p62dVTMVgRADQrWPkFLrozDN8KxbKMGqjVkPXBD7W");
     }
 }