Repository: mintlayer/mintlayer-core
Pull request #773 (2023-03-20 19:16:34)

Mine PoS blocks in tests
This is a fix for the problem with PoS unit tests discovered by Sam.
I added a mining function but only for `pos_processing tests`.  Not sure that it's useful elsewhere at this point.


Patch:
diff --git a/Cargo.lock b/Cargo.lock
index 36f046043..84f915998 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -643,6 +643,7 @@ version = "0.1.0"
 dependencies = [
  "chainstate-types",
  "common",
+ "crypto",
  "num",
  "pos_accounting",
  "rstest",
diff --git a/chainstate/src/detail/chainstateref/epoch_seal.rs b/chainstate/src/detail/chainstateref/epoch_seal.rs
index 3ac527109..94b4c66b3 100644
--- a/chainstate/src/detail/chainstateref/epoch_seal.rs
+++ b/chainstate/src/detail/chainstateref/epoch_seal.rs
@@ -173,7 +173,7 @@ fn create_randomness_from_block<S: BlockchainStorageRead>(
     let epoch_index = chain_config.epoch_index_from_height(block_height);
     PoSRandomness::from_block(
         epoch_index,
-        block.header(),
+        block.header().timestamp(),
         &sealed_epoch_randomness,
         pos_data,
         &vrf_pub_key,
diff --git a/chainstate/tx-verifier/src/transaction_verifier/mod.rs b/chainstate/tx-verifier/src/transaction_verifier/mod.rs
index 050ad936a..057f0e57d 100644
--- a/chainstate/tx-verifier/src/transaction_verifier/mod.rs
+++ b/chainstate/tx-verifier/src/transaction_verifier/mod.rs
@@ -28,7 +28,6 @@ pub mod flush;
 pub mod hierarchy;
 mod optional_tx_index_cache;
 pub mod storage;
-pub mod timelock_check;
 
 mod cached_operation;
 pub use cached_operation::CachedOperation;
@@ -52,7 +51,7 @@ use self::{
 };
 use ::utils::{ensure, shallow_clone::ShallowClone};
 
-use chainstate_types::BlockIndex;
+use chainstate_types::{block_index_ancestor_getter, BlockIndex, GenBlockIndex};
 use common::{
     amount_sum,
     chain::{
@@ -64,7 +63,7 @@ use common::{
         Block, ChainConfig, GenBlock, OutPointSourceId, OutputPurpose, Transaction, TxInput,
         TxMainChainIndex, TxOutput,
     },
-    primitives::{id::WithId, Amount, BlockHeight, Id, Idable, H256},
+    primitives::{id::WithId, Amount, BlockDistance, BlockHeight, Id, Idable, H256},
 };
 use consensus::ConsensusPoSError;
 use pos_accounting::{
@@ -486,6 +485,51 @@ where
         Ok(())
     }
 
+    fn check_timelock(
+        &self,
+        source_block_index: &GenBlockIndex,
+        output: &TxOutput,
+        spend_height: &BlockHeight,
+        spending_time: &BlockTimestamp,
+    ) -> Result<(), ConnectTransactionError> {
+        use common::chain::timelock::OutputTimeLock;
+
+        let timelock = match output.purpose() {
+            OutputPurpose::Transfer(_)
+            | OutputPurpose::Burn
+            | OutputPurpose::StakePool(_)
+            | OutputPurpose::ProduceBlockFromStake(_, _) => return Ok(()),
+            OutputPurpose::LockThenTransfer(_, tl) => tl,
+        };
+
+        let source_block_height = source_block_index.block_height();
+        let source_block_time = source_block_index.block_timestamp();
+
+        let past_lock = match timelock {
+            OutputTimeLock::UntilHeight(h) => spend_height >= h,
+            OutputTimeLock::UntilTime(t) => spending_time >= t,
+            OutputTimeLock::ForBlockCount(d) => {
+                let d: i64 = (*d)
+                    .try_into()
+                    .map_err(|_| ConnectTransactionError::BlockHeightArithmeticError)?;
+                let d = BlockDistance::from(d);
+                *spend_height
+                    >= (source_block_height + d)
+                        .ok_or(ConnectTransactionError::BlockHeightArithmeticError)?
+            }
+            OutputTimeLock::ForSeconds(dt) => {
+                *spending_time
+                    >= source_block_time
+                        .add_int_seconds(*dt)
+                        .ok_or(ConnectTransactionError::BlockTimestampArithmeticError)?
+            }
+        };
+
+        ensure!(past_lock, ConnectTransactionError::TimeLockViolation);
+
+        Ok(())
+    }
+
     fn verify_signatures<T: Transactable>(&self, tx: &T) -> Result<(), ConnectTransactionError> {
         let inputs = match tx.inputs() {
             Some(ins) => ins,
@@ -511,6 +555,72 @@ where
         Ok(())
     }
 
+    fn check_timelocks<T: Transactable>(
+        &self,
+        tx_source: &TransactionSourceForConnect,
+        tx: &T,
+        spending_time: &BlockTimestamp,
+    ) -> Result<(), ConnectTransactionError> {
+        let inputs = match tx.inputs() {
+            Some(ins) => ins,
+            None => return Ok(()),
+        };
+
+        for input in inputs {
+            let outpoint = input.outpoint();
+            let utxo = self
+                .utxo_cache
+                .utxo(outpoint)
+                .ok_or(ConnectTransactionError::MissingOutputOrSpent)?;
+
+            if utxo.output().has_timelock() {
+                let height = match utxo.source() {
+                    utxo::UtxoSource::Blockchain(h) => *h,
+                    utxo::UtxoSource::Mempool => match tx_source {
+                        TransactionSourceForConnect::Chain { new_block_index: _ } => {
+                            unreachable!("Mempool utxos can never be reached from storage while connecting local transactions")
+                        }
+                        TransactionSourceForConnect::Mempool { current_best } => {
+                            current_best.block_height().next_height()
+                        }
+                    },
+                };
+
+                let block_index_getter =
+                    |db_tx: &S, _chain_config: &ChainConfig, id: &Id<GenBlock>| {
+                        db_tx.get_gen_block_index(id)
+                    };
+
+                let starting_point: &BlockIndex = match tx_source {
+                    TransactionSourceForConnect::Chain { new_block_index } => new_block_index,
+                    TransactionSourceForConnect::Mempool { current_best } => current_best,
+                };
+
+                let source_block_index = block_index_ancestor_getter(
+                    block_index_getter,
+                    &self.storage,
+                    self.chain_config.as_ref(),
+                    (&starting_point.clone().into_gen_block_index()).into(),
+                    height,
+                )
+                .map_err(|e| {
+                    ConnectTransactionError::InvariantErrorHeaderCouldNotBeLoadedFromHeight(
+                        e, height,
+                    )
+                })?;
+
+                self.check_timelock(
+                    &source_block_index,
+                    utxo.output(),
+                    &tx_source.expected_block_height(),
+                    spending_time,
+                )?;
+            }
+        }
+
+        Ok(())
+    }
+
     fn connect_pos_accounting_outputs(
         &mut self,
         tx_source: TransactionSource,
@@ -630,14 +740,7 @@ where
         self.token_issuance_cache.register(block_id, tx.transaction())?;
 
         // check timelocks of the outputs and make sure there's no premature spending
-        timelock_check::check_timelocks(
-            &self.storage,
-            &self.chain_config,
-            &self.utxo_cache,
-            tx_source,
-            tx,
-            median_time_past,
-        )?;
+        self.check_timelocks(tx_source, tx, median_time_past)?;
 
         // verify input signatures
         self.verify_signatures(tx)?;
diff --git a/chainstate/tx-verifier/src/transaction_verifier/timelock_check.rs b/chainstate/tx-verifier/src/transaction_verifier/timelock_check.rs
deleted file mode 100644
index 05f779473..000000000
--- a/chainstate/tx-verifier/src/transaction_verifier/timelock_check.rs
+++ /dev/null
@@ -1,142 +0,0 @@
-// Copyright (c) 2023 RBB S.r.l
-// opensource@mintlayer.org
-// SPDX-License-Identifier: MIT
-// Licensed under the MIT License;
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-// https://github.com/mintlayer/mintlayer-core/blob/master/LICENSE
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-use chainstate_types::{block_index_ancestor_getter, BlockIndex, GenBlockIndex};
-use common::{
-    chain::{
-        block::timestamp::BlockTimestamp, signature::Transactable, ChainConfig, GenBlock,
-        OutputPurpose, TxOutput,
-    },
-    primitives::{BlockDistance, BlockHeight, Id},
-};
-use utils::ensure;
-use utxo::UtxosView;
-
-use super::{
-    error::ConnectTransactionError, storage::TransactionVerifierStorageRef,
-    TransactionSourceForConnect,
-};
-
-fn check_timelock(
-    source_block_index: &GenBlockIndex,
-    output: &TxOutput,
-    spend_height: &BlockHeight,
-    spending_time: &BlockTimestamp,
-) -> Result<(), ConnectTransactionError> {
-    use common::chain::timelock::OutputTimeLock;
-
-    let timelock = match output.purpose() {
-        OutputPurpose::Transfer(_)
-        | OutputPurpose::Burn
-        | OutputPurpose::StakePool(_)
-        | OutputPurpose::ProduceBlockFromStake(_, _) => return Ok(()),
-        OutputPurpose::LockThenTransfer(_, tl) => tl,
-    };
-
-    let source_block_height = source_block_index.block_height();
-    let source_block_time = source_block_index.block_timestamp();
-
-    let past_lock = match timelock {
-        OutputTimeLock::UntilHeight(h) => spend_height >= h,
-        OutputTimeLock::UntilTime(t) => spending_time >= t,
-        OutputTimeLock::ForBlockCount(d) => {
-            let d: i64 = (*d)
-                .try_into()
-                .map_err(|_| ConnectTransactionError::BlockHeightArithmeticError)?;
-            let d = BlockDistance::from(d);
-            *spend_height
-                >= (source_block_height + d)
-                    .ok_or(ConnectTransactionError::BlockHeightArithmeticError)?
-        }
-        OutputTimeLock::ForSeconds(dt) => {
-            *spending_time
-                >= source_block_time
-                    .add_int_seconds(*dt)
-                    .ok_or(ConnectTransactionError::BlockTimestampArithmeticError)?
-        }
-    };
-
-    ensure!(past_lock, ConnectTransactionError::TimeLockViolation);
-
-    Ok(())
-}
-
-pub fn check_timelocks<
-    S: TransactionVerifierStorageRef,
-    C: AsRef<ChainConfig>,
-    T: Transactable,
-    U: UtxosView,
->(
-    storage: &S,
-    chain_config: &C,
-    utxos_view: &U,
-    tx_source: &TransactionSourceForConnect,
-    tx: &T,
-    spending_time: &BlockTimestamp,
-) -> Result<(), ConnectTransactionError> {
-    let inputs = match tx.inputs() {
-        Some(ins) => ins,
-        None => return Ok(()),
-    };
-
-    for input in inputs {
-        let outpoint = input.outpoint();
-        let utxo =
-            utxos_view.utxo(outpoint).ok_or(ConnectTransactionError::MissingOutputOrSpent)?;
-
-        if utxo.output().has_timelock() {
-            let height = match utxo.source() {
-                utxo::UtxoSource::Blockchain(h) => *h,
-                utxo::UtxoSource::Mempool => match tx_source {
-                    TransactionSourceForConnect::Chain { new_block_index: _ } => {
-                        unreachable!("Mempool utxos can never be reached from storage while connecting local transactions")
-                    }
-                    TransactionSourceForConnect::Mempool { current_best } => {
-                        current_best.block_height().next_height()
-                    }
-                },
-            };
-
-            let block_index_getter = |db_tx: &S, _chain_config: &ChainConfig, id: &Id<GenBlock>| {
-                db_tx.get_gen_block_index(id)
-            };
-
-            let starting_point: &BlockIndex = match tx_source {
-                TransactionSourceForConnect::Chain { new_block_index } => new_block_index,
-                TransactionSourceForConnect::Mempool { current_best } => current_best,
-            };
-
-            let source_block_index = block_index_ancestor_getter(
-                block_index_getter,
-                storage,
-                chain_config.as_ref(),
-                (&starting_point.clone().into_gen_block_index()).into(),
-                height,
-            )
-            .map_err(|e| {
-                ConnectTransactionError::InvariantErrorHeaderCouldNotBeLoadedFromHeight(e, height)
-            })?;
-
-            check_timelock(
-                &source_block_index,
-                utxo.output(),
-                &tx_source.expected_block_height(),
-                spending_time,
-            )?;
-        }
-    }
-
-    Ok(())
-}
diff --git a/chainstate/types/src/pos_randomness.rs b/chainstate/types/src/pos_randomness.rs
index 01e2bcefb..f277b0a64 100644
--- a/chainstate/types/src/pos_randomness.rs
+++ b/chainstate/types/src/pos_randomness.rs
@@ -15,7 +15,7 @@
 
 use common::{
     chain::{
-        block::{consensus_data::PoSData, BlockHeader},
+        block::{consensus_data::PoSData, timestamp::BlockTimestamp},
         config::EpochIndex,
         Block, ChainConfig,
     },
@@ -47,7 +47,7 @@ impl PoSRandomness {
 
     pub fn from_block(
         epoch_index: EpochIndex,
-        header: &BlockHeader,
+        block_timestamp: BlockTimestamp,
         seal_randomness: &PoSRandomness,
         pos_data: &PoSData,
         vrf_pub_key: &VRFPublicKey,
@@ -57,7 +57,7 @@ impl PoSRandomness {
             &seal_randomness.value(),
             pos_data.vrf_data(),
             vrf_pub_key,
-            header,
+            block_timestamp,
         )?;
 
         Ok(Self::new(hash))
diff --git a/chainstate/types/src/vrf_tools.rs b/chainstate/types/src/vrf_tools.rs
index e9335e06f..5537e010e 100644
--- a/chainstate/types/src/vrf_tools.rs
+++ b/chainstate/types/src/vrf_tools.rs
@@ -14,10 +14,7 @@
 // limitations under the License.
 
 use common::{
-    chain::{
-        block::{timestamp::BlockTimestamp, BlockHeader},
-        config::EpochIndex,
-    },
+    chain::{block::timestamp::BlockTimestamp, config::EpochIndex},
     primitives::H256,
 };
 use crypto::vrf::{
@@ -78,10 +75,9 @@ pub fn verify_vrf_and_get_vrf_output(
     random_seed: &H256,
     vrf_data: &VRFReturn,
     vrf_public_key: &VRFPublicKey,
-    spender_block_header: &BlockHeader,
+    block_timestamp: BlockTimestamp,
 ) -> Result<H256, ProofOfStakeVRFError> {
-    let transcript =
-        construct_transcript(epoch_index, random_seed, spender_block_header.timestamp());
+    let transcript = construct_transcript(epoch_index, random_seed, block_timestamp);
 
     vrf_public_key.verify_vrf_data(transcript.clone().into(), vrf_data)?;
 
diff --git a/consensus/Cargo.toml b/consensus/Cargo.toml
index fa5c12a3f..22e5851b3 100644
--- a/consensus/Cargo.toml
+++ b/consensus/Cargo.toml
@@ -7,6 +7,7 @@ rust-version.workspace = true
 
 [dependencies]
 common = {path = '../common'}
+crypto = {path = '../crypto'}
 utils = {path = "../utils"}
 chainstate-types = {path = '../chainstate/types'}
 pos_accounting = {path = '../pos_accounting'}
diff --git a/consensus/src/lib.rs b/consensus/src/lib.rs
index c2df8f212..4abb8569d 100644
--- a/consensus/src/lib.rs
+++ b/consensus/src/lib.rs
@@ -20,7 +20,7 @@ mod pow;
 
 pub use crate::{
     error::ConsensusVerificationError,
-    pos::{error::ConsensusPoSError, kernel::get_kernel_output},
+    pos::{check_pos_hash, error::ConsensusPoSError, kernel::get_kernel_output},
     pow::{check_proof_of_work, mine, ConsensusPoWError},
     validator::validate_consensus,
 };
diff --git a/consensus/src/pos/mod.rs b/consensus/src/pos/mod.rs
index 04bca89fe..c324485b7 100644
--- a/consensus/src/pos/mod.rs
+++ b/consensus/src/pos/mod.rs
@@ -22,71 +22,47 @@ use chainstate_types::{
 };
 use common::{
     chain::{
-        block::{consensus_data::PoSData, BlockHeader},
+        block::{consensus_data::PoSData, timestamp::BlockTimestamp, BlockHeader},
         config::EpochIndex,
-        ChainConfig, OutputPurpose, TxOutput,
+        ChainConfig, OutputPurpose,
     },
-    primitives::{BlockHeight, Idable},
+    primitives::{Amount, BlockHeight, Idable},
     Uint256, Uint512,
 };
+use crypto::vrf::VRFPublicKey;
 use pos_accounting::PoSAccountingView;
 use utils::ensure;
 use utxo::UtxosView;
 
 use crate::pos::{error::ConsensusPoSError, kernel::get_kernel_output};
 
-fn check_stake_kernel_hash<P: PoSAccountingView>(
+pub fn check_pos_hash(
     epoch_index: EpochIndex,
     random_seed: &PoSRandomness,
     pos_data: &PoSData,
-    kernel_output: &TxOutput,
-    spender_block_header: &BlockHeader,
-    pos_accounting_view: &P,
+    vrf_pub_key: &VRFPublicKey,
+    block_timestamp: BlockTimestamp,
+    pool_balance: Amount,
 ) -> Result<(), ConsensusPoSError> {
     let target: Uint256 = (*pos_data.compact_target())
         .try_into()
         .map_err(|_| ConsensusPoSError::BitsToTargetConversionFailed(*pos_data.compact_target()))?;
 
-    let vrf_pub_key = match kernel_output.purpose() {
-        OutputPurpose::Transfer(_)
-        | OutputPurpose::LockThenTransfer(_, _)
-        | OutputPurpose::Burn => {
-            // only pool outputs can be staked
-            return Err(ConsensusPoSError::RandomnessError(
-                PoSRandomnessError::InvalidOutputPurposeInStakeKernel(
-                    spender_block_header.get_id(),
-                ),
-            ));
-        }
-        OutputPurpose::StakePool(d) => d.as_ref().vrf_public_key().clone(),
-        OutputPurpose::ProduceBlockFromStake(_, pool_id) => {
-            let pool_data = pos_accounting_view
-                .get_pool_data(*pool_id)?
-                .ok_or(ConsensusPoSError::PoolDataNotFound(*pool_id))?;
-            pool_data.vrf_public_key().clone()
-        }
-    };
-
-    let hash_pos: Uint256 = PoSRandomness::from_block(
+    let hash: Uint256 = PoSRandomness::from_block(
         epoch_index,
-        spender_block_header,
+        block_timestamp,
         random_seed,
         pos_data,
-        &vrf_pub_key,
+        vrf_pub_key,
     )?
     .value()
     .into();
 
-    let hash_pos_arith: Uint512 = hash_pos.into();
-
-    let stake_pool_id = *pos_data.stake_pool_id();
-    let pool_balance: Uint512 = pos_accounting_view
-        .get_pool_balance(stake_pool_id)?
-        .ok_or(ConsensusPoSError::PoolBalanceNotFound(stake_pool_id))?
-        .into();
+    let hash: Uint512 = hash.into();
+    let pool_balance: Uint512 = pool_balance.into();
 
     ensure!(
-        hash_pos_arith <= pool_balance * target.into(),
+        hash <= pool_balance * target.into(),
         ConsensusPoSError::StakeKernelHashTooHigh
     );
 
@@ -140,13 +116,37 @@ where
 
     let current_epoch_index = chain_config.epoch_index_from_height(&current_height);
     let kernel_output = get_kernel_output(pos_data.kernel_inputs(), utxos_view)?;
-    check_stake_kernel_hash(
+
+    let vrf_pub_key = match kernel_output.purpose() {
+        OutputPurpose::Transfer(_)
+        | OutputPurpose::LockThenTransfer(_, _)
+        | OutputPurpose::Burn => {
+            // only pool outputs can be staked
+            return Err(ConsensusPoSError::RandomnessError(
+                PoSRandomnessError::InvalidOutputPurposeInStakeKernel(header.get_id()),
+            ));
+        }
+        OutputPurpose::StakePool(d) => d.as_ref().vrf_public_key().clone(),
+        OutputPurpose::ProduceBlockFromStake(_, pool_id) => {
+            let pool_data = pos_accounting_view
+                .get_pool_data(*pool_id)?
+                .ok_or(ConsensusPoSError::PoolDataNotFound(*pool_id))?;
+            pool_data.vrf_public_key().clone()
+        }
+    };
+
+    let stake_pool_id = *pos_data.stake_pool_id();
+    let pool_balance = pos_accounting_view
+        .get_pool_balance(stake_pool_id)?
+        .ok_or(ConsensusPoSError::PoolBalanceNotFound(stake_pool_id))?;
+
+    check_pos_hash(
         current_epoch_index,
         &random_seed,
         pos_data,
-        &kernel_output,
-        header,
-        pos_accounting_view,
+        &vrf_pub_key,
+        header.timestamp(),
+        pool_balance,
     )?;
     Ok(())
 }


Test patch:
diff --git a/chainstate/test-suite/src/tests/pos_processing_tests.rs b/chainstate/test-suite/src/tests/pos_processing_tests.rs
index a14a59515..601784709 100644
--- a/chainstate/test-suite/src/tests/pos_processing_tests.rs
+++ b/chainstate/test-suite/src/tests/pos_processing_tests.rs
@@ -19,12 +19,12 @@ use chainstate::{
     chainstate_interface::ChainstateInterface, BlockError, ChainstateError, CheckBlockError,
     ConnectTransactionError, SpendStakeError,
 };
-use chainstate_storage::{BlockchainStorageRead, TipStorageTag, Transactional};
+use chainstate_storage::{BlockchainStorageRead, SealedStorageTag, TipStorageTag, Transactional};
 use chainstate_test_framework::{
     anyonecanspend_address, empty_witness, TestFramework, TransactionBuilder,
 };
 use chainstate_types::{
-    pos_randomness::PoSRandomnessError,
+    pos_randomness::{PoSRandomness, PoSRandomnessError},
     vrf_tools::{construct_transcript, ProofOfStakeVRFError},
 };
 use common::{
@@ -41,7 +41,7 @@ use common::{
     Uint256,
 };
 use consensus::{ConsensusPoSError, ConsensusVerificationError};
-use crypto::random::CryptoRng;
+use crypto::{random::CryptoRng, vrf::VRFPublicKey};
 use crypto::{
     random::Rng,
     vrf::{VRFError, VRFKeyKind, VRFPrivateKey},
@@ -58,18 +58,14 @@ const TEST_EPOCH_LENGTH: NonZeroU64 = match NonZeroU64::new(2) {
 };
 const TEST_SEALED_EPOCH_DISTANCE: usize = 0;
 
-fn create_stake_pool_data(rng: &mut (impl Rng + CryptoRng)) -> (VRFPrivateKey, StakePoolData) {
+fn create_stake_pool_data(rng: &mut (impl Rng + CryptoRng), vrf_pk: VRFPublicKey) -> StakePoolData {
     let destination = super::new_pub_key_destination(rng);
-    let (vrf_sk, vrf_pk) = VRFPrivateKey::new_from_rng(rng, VRFKeyKind::Schnorrkel);
-    (
-        vrf_sk,
-        StakePoolData::new(
-            anyonecanspend_address(),
-            vrf_pk,
-            destination,
-            0,
-            Amount::ZERO,
-        ),
+    StakePoolData::new(
+        anyonecanspend_address(),
+        vrf_pk,
+        destination,
+        0,
+        Amount::ZERO,
     )
 }
 
@@ -152,40 +148,56 @@ fn add_block_with_2_stake_pools(
     (stake_outpoint1, pool_id1, outpoint2, pool_id2)
 }
 
-fn create_pos_data(
-    tf: &mut TestFramework,
-    outpoint: OutPoint,
+#[allow(clippy::too_many_arguments)]
+fn mine(
+    initial_timestamp: BlockTimestamp,
+    kernel_outpoint: OutPoint,
+    vrf_pk: &VRFPublicKey,
     vrf_sk: &VRFPrivateKey,
-    sealed_epoch_randomness: H256,
+    sealed_epoch_randomness: PoSRandomness,
     pool_id: PoolId,
+    pool_balance: Amount,
     epoch_index: EpochIndex,
-) -> PoSData {
+) -> Option<(PoSData, BlockTimestamp)> {
     let difficulty = Uint256::MAX;
-    let block_timestamp = BlockTimestamp::from_duration_since_epoch(tf.current_time());
+    let mut timestamp = initial_timestamp;
 
-    let vrf_sealed_epoch_transcript =
-        construct_transcript(epoch_index, &sealed_epoch_randomness, block_timestamp);
-    let vrf_data_from_sealed_epoch = vrf_sk.produce_vrf_data(vrf_sealed_epoch_transcript.into());
+    for _ in 0..1000 {
+        let transcript =
+            construct_transcript(epoch_index, &sealed_epoch_randomness.value(), timestamp);
+        let vrf_data = vrf_sk.produce_vrf_data(transcript.into());
+        let pos_data = PoSData::new(
+            vec![kernel_outpoint.clone().into()],
+            vec![InputWitness::NoSignature(None)],
+            pool_id,
+            vrf_data,
+            Compact::from(difficulty),
+        );
 
-    PoSData::new(
-        vec![outpoint.into()],
-        vec![InputWitness::NoSignature(None)],
-        pool_id,
-        vrf_data_from_sealed_epoch,
-        Compact::from(difficulty),
-    )
-}
+        if consensus::check_pos_hash(
+            epoch_index,
+            &sealed_epoch_randomness,
+            &pos_data,
+            vrf_pk,
+            timestamp,
+            pool_balance,
+        )
+        .is_ok()
+        {
+            return Some((pos_data, timestamp));
+        }
 
-// Create a chain genesis <- block_1 <- block_2
-// PoS consensus activates on height 2.
-// block_1 has valid StakePool output. block_2 has PoS kernel input from block_1.
-// Check that the chain is valid.
-#[rstest]
-#[trace]
-#[case(Seed::from_entropy())]
-fn pos_basic(#[case] seed: Seed) {
-    let mut rng = make_seedable_rng(seed);
+        timestamp = timestamp.add_int_seconds(1).unwrap();
+    }
+    None
+}
 
+// Create a chain genesis <- block_1
+// block_1 has tx with StakePool output
+fn setup_test_chain_with_staked_pool(
+    rng: &mut (impl Rng + CryptoRng),
+    vrf_pk: VRFPublicKey,
+) -> (TestFramework, OutPoint, PoolId) {
     let upgrades = vec![
         (
             BlockHeight::new(0),
@@ -202,27 +214,51 @@ fn pos_basic(#[case] seed: Seed) {
         .epoch_length(TEST_EPOCH_LENGTH)
         .sealed_epoch_distance_from_tip(TEST_SEALED_EPOCH_DISTANCE)
         .build();
-    let mut tf = TestFramework::builder(&mut rng).with_chain_config(chain_config).build();
 
-    let (vrf_sk, stake_pool_data) = create_stake_pool_data(&mut rng);
-    let (stake_pool_outpoint, pool_id) =
-        add_block_with_stake_pool(&mut rng, &mut tf, stake_pool_data);
+    let mut tf = TestFramework::builder(rng).with_chain_config(chain_config).build();
+
+    let stake_pool_data = create_stake_pool_data(rng, vrf_pk);
+    let (stake_pool_outpoint, pool_id) = add_block_with_stake_pool(rng, &mut tf, stake_pool_data);
+
+    (tf, stake_pool_outpoint, pool_id)
+}
+
+// Create a chain genesis <- block_1 <- block_2
+// PoS consensus activates on height 2.
+// block_1 has valid StakePool output. block_2 has PoS kernel input from block_1.
+// Check that the chain is valid.
+#[rstest]
+#[trace]
+#[case(Seed::from_entropy())]
+fn pos_basic(#[case] seed: Seed) {
+    let mut rng = make_seedable_rng(seed);
+    let (vrf_sk, vrf_pk) = VRFPrivateKey::new_from_rng(&mut rng, VRFKeyKind::Schnorrkel);
+    let (mut tf, stake_pool_outpoint, pool_id) =
+        setup_test_chain_with_staked_pool(&mut rng, vrf_pk.clone());
 
     let initial_randomness = tf.chainstate.get_chain_config().initial_randomness();
-    let pos_data = create_pos_data(
-        &mut tf,
+    let sealed_pool_balance =
+        PoSAccountingStorageRead::<SealedStorageTag>::get_pool_balance(&tf.storage, pool_id)
+            .unwrap()
+            .unwrap();
+    let (pos_data, block_timestamp) = mine(
+        BlockTimestamp::from_duration_since_epoch(tf.current_time()),
         stake_pool_outpoint,
+        &vrf_pk,
         &vrf_sk,
-        initial_randomness,
+        PoSRandomness::new(initial_randomness),
         pool_id,
+        sealed_pool_balance,
         1,
-    );
+    )
+    .expect("should be able to mine");
     let consensus_data = ConsensusData::PoS(Box::new(pos_data));
 
     // skip block reward output
     let res = tf
         .make_block_builder()
         .with_consensus_data(consensus_data.clone())
+        .with_timestamp(block_timestamp)
         .build_and_process()
         .unwrap_err();
     assert_eq!(
@@ -239,6 +275,7 @@ fn pos_basic(#[case] seed: Seed) {
     );
     tf.make_block_builder()
         .with_consensus_data(consensus_data)
+        .with_timestamp(block_timestamp)
         .with_reward(vec![reward_output])
         .build_and_process()
         .unwrap();
@@ -280,25 +317,30 @@ fn pos_invalid_kernel_input(#[case] seed: Seed) {
     let mut tf = TestFramework::builder(&mut rng).with_chain_config(chain_config).build();
 
     let genesis_id = tf.genesis().get_id();
-    let (vrf_sk, _) = VRFPrivateKey::new_from_rng(&mut rng, VRFKeyKind::Schnorrkel);
+    let (vrf_sk, vrf_pk) = VRFPrivateKey::new_from_rng(&mut rng, VRFKeyKind::Schnorrkel);
     let pool_id = pos_accounting::make_pool_id(&OutPoint::new(
         OutPointSourceId::BlockReward(genesis_id.into()),
         0,
     ));
 
+    let invalid_kernel_input = OutPoint::new(OutPointSourceId::BlockReward(genesis_id.into()), 0);
     let initial_randomness = tf.chainstate.get_chain_config().initial_randomness();
-    let pos_data = create_pos_data(
-        &mut tf,
-        OutPoint::new(OutPointSourceId::BlockReward(genesis_id.into()), 0),
+    let (pos_data, block_timestamp) = mine(
+        BlockTimestamp::from_duration_since_epoch(tf.current_time()),
+        invalid_kernel_input,
+        &vrf_pk,
         &vrf_sk,
-        initial_randomness,
+        PoSRandomness::new(initial_randomness),
         pool_id,
+        Amount::from_atoms(1),
         1,
-    );
+    )
+    .expect("should be able to mine");
 
     let res = tf
         .make_block_builder()
         .with_consensus_data(ConsensusData::PoS(Box::new(pos_data)))
+        .with_timestamp(block_timestamp)
         .build_and_process()
         .unwrap_err();
 
@@ -324,29 +366,9 @@ fn pos_invalid_kernel_input(#[case] seed: Seed) {
 #[case(Seed::from_entropy())]
 fn pos_invalid_vrf(#[case] seed: Seed) {
     let mut rng = make_seedable_rng(seed);
-
-    let difficulty = Uint256::MAX;
-    let upgrades = vec![
-        (
-            BlockHeight::new(0),
-            UpgradeVersion::ConsensusUpgrade(ConsensusUpgrade::IgnoreConsensus),
-        ),
-        (
-            BlockHeight::new(2),
-            UpgradeVersion::ConsensusUpgrade(ConsensusUpgrade::PoS),
-        ),
-    ];
-    let net_upgrades = NetUpgrades::initialize(upgrades).expect("valid net-upgrades");
-    let chain_config = ConfigBuilder::test_chain()
-        .net_upgrades(net_upgrades)
-        .epoch_length(TEST_EPOCH_LENGTH)
-        .sealed_epoch_distance_from_tip(TEST_SEALED_EPOCH_DISTANCE)
-        .build();
-    let mut tf = TestFramework::builder(&mut rng).with_chain_config(chain_config).build();
-
-    let (vrf_sk, stake_pool_data) = create_stake_pool_data(&mut rng);
-    let (stake_pool_outpoint, pool_id) =
-        add_block_with_stake_pool(&mut rng, &mut tf, stake_pool_data);
+    let (vrf_sk, vrf_pk) = VRFPrivateKey::new_from_rng(&mut rng, VRFKeyKind::Schnorrkel);
+    let (mut tf, stake_pool_outpoint, pool_id) =
+        setup_test_chain_with_staked_pool(&mut rng, vrf_pk.clone());
 
     let expected_error = ChainstateError::ProcessBlockError(BlockError::CheckBlockFailed(
         CheckBlockError::ConsensusVerificationFailed(ConsensusVerificationError::PoSError(
@@ -361,7 +383,22 @@ fn pos_invalid_vrf(#[case] seed: Seed) {
     let valid_epoch: EpochIndex = 1;
     let valid_vrf_transcript =
         construct_transcript(valid_epoch, &valid_prev_randomness, valid_block_timestamp);
-    let valid_vrf_data = vrf_sk.produce_vrf_data(valid_vrf_transcript.clone().into());
+
+    let sealed_pool_balance =
+        PoSAccountingStorageRead::<SealedStorageTag>::get_pool_balance(&tf.storage, pool_id)
+            .unwrap()
+            .unwrap();
+    let (valid_pos_data, valid_block_timestamp) = mine(
+        BlockTimestamp::from_duration_since_epoch(tf.current_time()),
+        stake_pool_outpoint,
+        &vrf_pk,
+        &vrf_sk,
+        PoSRandomness::new(valid_prev_randomness),
+        pool_id,
+        sealed_pool_balance,
+        valid_epoch,
+    )
+    .expect("should be able to mine");
 
     {
         // invalid sealed epoch randomness
@@ -370,11 +407,11 @@ fn pos_invalid_vrf(#[case] seed: Seed) {
             construct_transcript(valid_epoch, &invalid_randomness, valid_block_timestamp);
         let vrf_data = vrf_sk.produce_vrf_data(vrf_transcript.into());
         let pos_data = PoSData::new(
-            vec![stake_pool_outpoint.clone().into()],
-            vec![InputWitness::NoSignature(None)],
+            valid_pos_data.kernel_inputs().to_owned(),
+            valid_pos_data.kernel_witness().to_owned(),
             pool_id,
             vrf_data,
-            Compact::from(difficulty),
+            *valid_pos_data.compact_target(),
         );
 
         let res = tf
@@ -388,16 +425,16 @@ fn pos_invalid_vrf(#[case] seed: Seed) {
 
     {
         // invalid timestamp
-        let block_timestamp =
-            BlockTimestamp::from_duration_since_epoch(tf.current_time().saturating_mul(2));
-        let vrf_transcript = construct_transcript(1, &valid_prev_randomness, block_timestamp);
+        let block_timestamp = valid_block_timestamp.add_int_seconds(1).unwrap();
+        let vrf_transcript =
+            construct_transcript(valid_epoch, &valid_prev_randomness, block_timestamp);
         let vrf_data = vrf_sk.produce_vrf_data(vrf_transcript.into());
         let pos_data = PoSData::new(
-            vec![stake_pool_outpoint.clone().into()],
-            vec![InputWitness::NoSignature(None)],
+            valid_pos_data.kernel_inputs().to_owned(),
+            valid_pos_data.kernel_witness().to_owned(),
             pool_id,
             vrf_data,
-            Compact::from(difficulty),
+            *valid_pos_data.compact_target(),
         );
 
         let res = tf
@@ -411,14 +448,18 @@ fn pos_invalid_vrf(#[case] seed: Seed) {
 
     {
         // invalid epoch
-        let vrf_transcript = construct_transcript(2, &valid_prev_randomness, valid_block_timestamp);
+        let vrf_transcript = construct_transcript(
+            valid_epoch + 1,
+            &valid_prev_randomness,
+            valid_block_timestamp,
+        );
         let vrf_data = vrf_sk.produce_vrf_data(vrf_transcript.into());
         let pos_data = PoSData::new(
-            vec![stake_pool_outpoint.clone().into()],
-            vec![InputWitness::NoSignature(None)],
+            valid_pos_data.kernel_inputs().to_owned(),
+            valid_pos_data.kernel_witness().to_owned(),
             pool_id,
             vrf_data,
-            Compact::from(difficulty),
+            *valid_pos_data.compact_target(),
         );
 
         let res = tf
@@ -435,11 +476,11 @@ fn pos_invalid_vrf(#[case] seed: Seed) {
         let (vrf_sk_2, _) = VRFPrivateKey::new_from_rng(&mut rng, VRFKeyKind::Schnorrkel);
         let vrf_data = vrf_sk_2.produce_vrf_data(valid_vrf_transcript.into());
         let pos_data = PoSData::new(
-            vec![stake_pool_outpoint.clone().into()],
-            vec![InputWitness::NoSignature(None)],
+            valid_pos_data.kernel_inputs().to_owned(),
+            valid_pos_data.kernel_witness().to_owned(),
             pool_id,
             vrf_data,
-            Compact::from(difficulty),
+            *valid_pos_data.compact_target(),
         );
 
         let res = tf
@@ -453,15 +494,7 @@ fn pos_invalid_vrf(#[case] seed: Seed) {
 
     {
         // valid case
-        let pos_data = PoSData::new(
-            vec![stake_pool_outpoint.into()],
-            vec![InputWitness::NoSignature(None)],
-            pool_id,
-            valid_vrf_data,
-            Compact::from(difficulty),
-        );
-
-        let consensus_data = ConsensusData::PoS(Box::new(pos_data));
+        let consensus_data = ConsensusData::PoS(Box::new(valid_pos_data));
         let reward_output = TxOutput::new(
             OutputValue::Coin(Amount::from_atoms(1)),
             OutputPurpose::ProduceBlockFromStake(anyonecanspend_address(), pool_id),
@@ -469,6 +502,7 @@ fn pos_invalid_vrf(#[case] seed: Seed) {
         tf.make_block_builder()
             .with_consensus_data(consensus_data)
             .with_reward(vec![reward_output])
+            .with_timestamp(valid_block_timestamp)
             .build_and_process()
             .unwrap();
     }
@@ -483,43 +517,39 @@ fn pos_invalid_vrf(#[case] seed: Seed) {
 #[case(Seed::from_entropy())]
 fn pos_invalid_pool_id(#[case] seed: Seed) {
     let mut rng = make_seedable_rng(seed);
+    let (vrf_sk, vrf_pk) = VRFPrivateKey::new_from_rng(&mut rng, VRFKeyKind::Schnorrkel);
+    let (mut tf, stake_pool_outpoint, pool_id) =
+        setup_test_chain_with_staked_pool(&mut rng, vrf_pk.clone());
 
-    let upgrades = vec![
-        (
-            BlockHeight::new(0),
-            UpgradeVersion::ConsensusUpgrade(ConsensusUpgrade::IgnoreConsensus),
-        ),
-        (
-            BlockHeight::new(2),
-            UpgradeVersion::ConsensusUpgrade(ConsensusUpgrade::PoS),
-        ),
-    ];
-    let net_upgrades = NetUpgrades::initialize(upgrades).expect("valid net-upgrades");
-    let chain_config = ConfigBuilder::test_chain()
-        .net_upgrades(net_upgrades)
-        .epoch_length(TEST_EPOCH_LENGTH)
-        .sealed_epoch_distance_from_tip(TEST_SEALED_EPOCH_DISTANCE)
-        .build();
-    let mut tf = TestFramework::builder(&mut rng).with_chain_config(chain_config).build();
-
-    let (vrf_sk, stake_pool_data) = create_stake_pool_data(&mut rng);
-    let (stake_pool_outpoint, expected_pool_id) =
-        add_block_with_stake_pool(&mut rng, &mut tf, stake_pool_data);
-
-    let random_pool_id: PoolId = H256::random_using(&mut rng).into();
     let initial_randomness = tf.chainstate.get_chain_config().initial_randomness();
-    let pos_data = create_pos_data(
-        &mut tf,
-        stake_pool_outpoint.clone(),
+    let sealed_pool_balance =
+        PoSAccountingStorageRead::<SealedStorageTag>::get_pool_balance(&tf.storage, pool_id)
+            .unwrap()
+            .unwrap();
+    let (valid_pos_data, block_timestamp) = mine(
+        BlockTimestamp::from_duration_since_epoch(tf.current_time()),
+        stake_pool_outpoint,
+        &vrf_pk,
         &vrf_sk,
-        initial_randomness,
-        random_pool_id,
+        PoSRandomness::new(initial_randomness),
+        pool_id,
+        sealed_pool_balance,
         1,
+    )
+    .expect("should be able to mine");
+
+    let random_pool_id: PoolId = H256::random_using(&mut rng).into();
+    let invalid_pos_data = PoSData::new(
+        valid_pos_data.kernel_inputs().to_owned(),
+        valid_pos_data.kernel_witness().to_owned(),
+        random_pool_id,
+        valid_pos_data.vrf_data().clone(),
+        *valid_pos_data.compact_target(),
     );
 
     let res = tf
         .make_block_builder()
-        .with_consensus_data(ConsensusData::PoS(Box::new(pos_data)))
+        .with_consensus_data(ConsensusData::PoS(Box::new(invalid_pos_data)))
         .build_and_process()
         .unwrap_err();
 
@@ -533,22 +563,14 @@ fn pos_invalid_pool_id(#[case] seed: Seed) {
     );
 
     // test valid case
-    let pos_data = create_pos_data(
-        &mut tf,
-        stake_pool_outpoint,
-        &vrf_sk,
-        initial_randomness,
-        expected_pool_id,
-        1,
-    );
-
     let reward_output = TxOutput::new(
         OutputValue::Coin(Amount::from_atoms(1)),
-        OutputPurpose::ProduceBlockFromStake(anyonecanspend_address(), expected_pool_id),
+        OutputPurpose::ProduceBlockFromStake(anyonecanspend_address(), pool_id),
     );
     tf.make_block_builder()
-        .with_consensus_data(ConsensusData::PoS(Box::new(pos_data)))
+        .with_consensus_data(ConsensusData::PoS(Box::new(valid_pos_data)))
         .with_reward(vec![reward_output])
+        .with_timestamp(block_timestamp)
         .build_and_process()
         .unwrap();
 }
@@ -561,6 +583,7 @@ fn pos_invalid_pool_id(#[case] seed: Seed) {
 #[case(Seed::from_entropy())]
 fn not_sealed_pool_cannot_be_used(#[case] seed: Seed) {
     let mut rng = make_seedable_rng(seed);
+    let (vrf_sk, vrf_pk) = VRFPrivateKey::new_from_rng(&mut rng, VRFKeyKind::Schnorrkel);
 
     let upgrades = vec![
         (
@@ -580,23 +603,27 @@ fn not_sealed_pool_cannot_be_used(#[case] seed: Seed) {
         .build();
     let mut tf = TestFramework::builder(&mut rng).with_chain_config(chain_config).build();
 
-    let (vrf_sk, stake_pool_data) = create_stake_pool_data(&mut rng);
+    let stake_pool_data = create_stake_pool_data(&mut rng, vrf_pk.clone());
     let (stake_pool_outpoint, pool_id) =
         add_block_with_stake_pool(&mut rng, &mut tf, stake_pool_data);
 
     let initial_randomness = tf.chainstate.get_chain_config().initial_randomness();
-    let pos_data = create_pos_data(
-        &mut tf,
+    let (pos_data, block_timestamp) = mine(
+        BlockTimestamp::from_duration_since_epoch(tf.current_time()),
         stake_pool_outpoint,
+        &vrf_pk,
         &vrf_sk,
-        initial_randomness,
+        PoSRandomness::new(initial_randomness),
         pool_id,
-        0,
-    );
+        Amount::from_atoms(1),
+        1,
+    )
+    .expect("should be able to mine");
 
     let res = tf
         .make_block_builder()
         .with_consensus_data(ConsensusData::PoS(Box::new(pos_data)))
+        .with_timestamp(block_timestamp)
         .build_and_process()
         .unwrap_err();
 
@@ -627,41 +654,24 @@ fn not_sealed_pool_cannot_be_used(#[case] seed: Seed) {
 #[case(Seed::from_entropy())]
 fn spend_stake_pool_in_block_reward(#[case] seed: Seed) {
     let mut rng = make_seedable_rng(seed);
-
-    let upgrades = vec![
-        (
-            BlockHeight::new(0),
-            UpgradeVersion::ConsensusUpgrade(ConsensusUpgrade::IgnoreConsensus),
-        ),
-        (
-            BlockHeight::new(2),
-            UpgradeVersion::ConsensusUpgrade(ConsensusUpgrade::PoS),
-        ),
-    ];
-    let net_upgrades = NetUpgrades::initialize(upgrades).expect("valid net-upgrades");
-    let chain_config = ConfigBuilder::test_chain()
-        .net_upgrades(net_upgrades)
-        .epoch_length(TEST_EPOCH_LENGTH)
-        .sealed_epoch_distance_from_tip(TEST_SEALED_EPOCH_DISTANCE)
-        .build();
-    let mut tf = TestFramework::builder(&mut rng).with_chain_config(chain_config).build();
-
-    // create initial chain: genesis <- block_1
-    let (vrf_sk, stake_pool_data) = create_stake_pool_data(&mut rng);
-    let (stake_pool_outpoint, pool_id) =
-        add_block_with_stake_pool(&mut rng, &mut tf, stake_pool_data);
+    let (vrf_sk, vrf_pk) = VRFPrivateKey::new_from_rng(&mut rng, VRFKeyKind::Schnorrkel);
+    let (mut tf, stake_pool_outpoint, pool_id) =
+        setup_test_chain_with_staked_pool(&mut rng, vrf_pk.clone());
 
     // prepare and process block_2 with StakePool -> ProduceBlockFromStake kernel
     let initial_randomness = tf.chainstate.get_chain_config().initial_randomness();
-    let pos_data = create_pos_data(
-        &mut tf,
+    let (pos_data, block_timestamp) = mine(
+        BlockTimestamp::from_duration_since_epoch(tf.current_time()),
         stake_pool_outpoint,
+        &vrf_pk,
         &vrf_sk,
         // no epoch is sealed yet so use initial randomness
-        initial_randomness,
+        PoSRandomness::new(initial_randomness),
         pool_id,
+        Amount::from_atoms(1),
         1,
-    );
+    )
+    .expect("should be able to mine");
     let reward_output = TxOutput::new(
         OutputValue::Coin(Amount::from_atoms(1)),
         OutputPurpose::ProduceBlockFromStake(anyonecanspend_address(), pool_id),
@@ -669,6 +679,7 @@ fn spend_stake_pool_in_block_reward(#[case] seed: Seed) {
     tf.make_block_builder()
         .with_consensus_data(ConsensusData::PoS(Box::new(pos_data)))
         .with_reward(vec![reward_output])
+        .with_timestamp(block_timestamp)
         .build_and_process()
         .unwrap();
 
@@ -677,15 +688,18 @@ fn spend_stake_pool_in_block_reward(#[case] seed: Seed) {
         OutPointSourceId::BlockReward(tf.chainstate.get_best_block_id().unwrap()),
         0,
     );
-    let pos_data = create_pos_data(
-        &mut tf,
+    let (pos_data, block_timestamp) = mine(
+        block_timestamp,
         block_2_reward_outpoint,
+        &vrf_pk,
         &vrf_sk,
         // no epoch is sealed yet so use initial randomness
-        initial_randomness,
+        PoSRandomness::new(initial_randomness),
         pool_id,
+        Amount::from_atoms(1),
         1,
-    );
+    )
+    .expect("should be able to mine");
     let reward_output = TxOutput::new(
         OutputValue::Coin(Amount::from_atoms(1)),
         OutputPurpose::ProduceBlockFromStake(anyonecanspend_address(), pool_id),
@@ -693,6 +707,7 @@ fn spend_stake_pool_in_block_reward(#[case] seed: Seed) {
     tf.make_block_builder()
         .with_consensus_data(ConsensusData::PoS(Box::new(pos_data)))
         .with_reward(vec![reward_output])
+        .with_timestamp(block_timestamp)
         .build_and_process()
         .unwrap();
 
@@ -705,14 +720,21 @@ fn spend_stake_pool_in_block_reward(#[case] seed: Seed) {
     // both sealed epoch and pre block randomness can be used
     let sealed_epoch_randomness =
         tf.storage.transaction_ro().unwrap().get_epoch_data(1).unwrap().unwrap();
-    let pos_data = create_pos_data(
-        &mut tf,
+    let sealed_pool_balance =
+        PoSAccountingStorageRead::<SealedStorageTag>::get_pool_balance(&tf.storage, pool_id)
+            .unwrap()
+            .unwrap();
+    let (pos_data, block_timestamp) = mine(
+        block_timestamp,
         block_3_reward_outpoint,
+        &vrf_pk,
         &vrf_sk,
-        sealed_epoch_randomness.randomness().value(),
+        *sealed_epoch_randomness.randomness(),
         pool_id,
+        sealed_pool_balance,
         2,
-    );
+    )
+    .expect("should be able to mine");
     let reward_output = TxOutput::new(
         OutputValue::Coin(Amount::from_atoms(1)),
         OutputPurpose::ProduceBlockFromStake(anyonecanspend_address(), pool_id),
@@ -720,6 +742,7 @@ fn spend_stake_pool_in_block_reward(#[case] seed: Seed) {
     tf.make_block_builder()
         .with_consensus_data(ConsensusData::PoS(Box::new(pos_data)))
         .with_reward(vec![reward_output])
+        .with_timestamp(block_timestamp)
         .build_and_process()
         .unwrap();
 
@@ -741,6 +764,8 @@ fn spend_stake_pool_in_block_reward(#[case] seed: Seed) {
 #[case(Seed::from_entropy())]
 fn mismatched_pools_in_kernel_and_reward(#[case] seed: Seed) {
     let mut rng = make_seedable_rng(seed);
+    let (vrf_sk_1, vrf_pk_1) = VRFPrivateKey::new_from_rng(&mut rng, VRFKeyKind::Schnorrkel);
+    let (_, vrf_pk_2) = VRFPrivateKey::new_from_rng(&mut rng, VRFKeyKind::Schnorrkel);
 
     let upgrades = vec![
         (
@@ -762,23 +787,29 @@ fn mismatched_pools_in_kernel_and_reward(#[case] seed: Seed) {
 
     // create initial chain: genesis <- block_1
     // block1 creates 2 separate pools
-    let (vrf_sk, stake_pool_data1) = create_stake_pool_data(&mut rng);
-    let (_, stake_pool_data2) = create_stake_pool_data(&mut rng);
+    let stake_pool_data1 = create_stake_pool_data(&mut rng, vrf_pk_1.clone());
+    let stake_pool_data2 = create_stake_pool_data(&mut rng, vrf_pk_2);
     let (stake_pool_outpoint1, pool_id1, _, pool_id2) =
         add_block_with_2_stake_pools(&mut rng, &mut tf, stake_pool_data1, stake_pool_data2);
 
     // prepare and process block_2 with StakePool -> ProduceBlockFromStake kernel
     // kernel refers to pool1, while block reward refers to pool2
     let initial_randomness = tf.chainstate.get_chain_config().initial_randomness();
-    let pos_data = create_pos_data(
-        &mut tf,
+    let sealed_pool_balance =
+        PoSAccountingStorageRead::<SealedStorageTag>::get_pool_balance(&tf.storage, pool_id1)
+            .unwrap()
+            .unwrap();
+    let (pos_data, block_timestamp) = mine(
+        BlockTimestamp::from_duration_since_epoch(tf.current_time()),
         stake_pool_outpoint1,
-        &vrf_sk,
-        // no epoch is sealed yet so use initial randomness
-        initial_randomness,
+        &vrf_pk_1,
+        &vrf_sk_1,
+        PoSRandomness::new(initial_randomness),
         pool_id1,
+        sealed_pool_balance,
         1,
-    );
+    )
+    .expect("should be able to mine");
     let reward_output = TxOutput::new(
         OutputValue::Coin(Amount::from_atoms(1)),
         OutputPurpose::ProduceBlockFromStake(anyonecanspend_address(), pool_id2),
@@ -787,6 +818,7 @@ fn mismatched_pools_in_kernel_and_reward(#[case] seed: Seed) {
         .make_block_builder()
         .with_consensus_data(ConsensusData::PoS(Box::new(pos_data)))
         .with_reward(vec![reward_output])
+        .with_timestamp(block_timestamp)
         .build_and_process()
         .unwrap_err();
 
@@ -804,48 +836,20 @@ fn mismatched_pools_in_kernel_and_reward(#[case] seed: Seed) {
 fn stake_pool_as_reward_output(#[case] seed: Seed) {
     let mut rng = make_seedable_rng(seed);
 
-    let upgrades = vec![
-        (
-            BlockHeight::new(0),
-            UpgradeVersion::ConsensusUpgrade(ConsensusUpgrade::IgnoreConsensus),
-        ),
-        (
-            BlockHeight::new(2),
-            UpgradeVersion::ConsensusUpgrade(ConsensusUpgrade::PoS),
-        ),
-    ];
-    let net_upgrades = NetUpgrades::initialize(upgrades).expect("valid net-upgrades");
     let chain_config = ConfigBuilder::test_chain()
-        .net_upgrades(net_upgrades)
         .epoch_length(TEST_EPOCH_LENGTH)
         .sealed_epoch_distance_from_tip(TEST_SEALED_EPOCH_DISTANCE)
         .build();
     let mut tf = TestFramework::builder(&mut rng).with_chain_config(chain_config).build();
 
-    let (vrf_sk, stake_pool_data) = create_stake_pool_data(&mut rng);
-    let (stake_pool_outpoint, pool_id) =
-        add_block_with_stake_pool(&mut rng, &mut tf, stake_pool_data.clone());
+    let (_, vrf_pk) = VRFPrivateKey::new_from_rng(&mut rng, VRFKeyKind::Schnorrkel);
+    let stake_pool_data = create_stake_pool_data(&mut rng, vrf_pk);
 
-    let initial_randomness = tf.chainstate.get_chain_config().initial_randomness();
-    let pos_data = create_pos_data(
-        &mut tf,
-        stake_pool_outpoint,
-        &vrf_sk,
-        initial_randomness,
-        pool_id,
-        1,
-    );
-
-    // valid case
     let reward_output = TxOutput::new(
         OutputValue::Coin(Amount::from_atoms(1)),
         OutputPurpose::StakePool(Box::new(stake_pool_data)),
     );
-    let block = tf
-        .make_block_builder()
-        .with_consensus_data(ConsensusData::PoS(Box::new(pos_data)))
-        .with_reward(vec![reward_output])
-        .build();
+    let block = tf.make_block_builder().with_reward(vec![reward_output]).build();
     let block_id = block.get_id();
     assert_eq!(
         tf.process_block(block, chainstate::BlockSource::Local).unwrap_err(),
@@ -867,42 +871,31 @@ fn stake_pool_as_reward_output(#[case] seed: Seed) {
 #[case(Seed::from_entropy())]
 fn check_pool_balance_after_reorg(#[case] seed: Seed) {
     let mut rng = make_seedable_rng(seed);
-
-    let upgrades = vec![
-        (
-            BlockHeight::new(0),
-            UpgradeVersion::ConsensusUpgrade(ConsensusUpgrade::IgnoreConsensus),
-        ),
-        (
-            BlockHeight::new(2),
-            UpgradeVersion::ConsensusUpgrade(ConsensusUpgrade::PoS),
-        ),
-    ];
-    let net_upgrades = NetUpgrades::initialize(upgrades).expect("valid net-upgrades");
-    let chain_config = ConfigBuilder::test_chain()
-        .net_upgrades(net_upgrades)
-        .epoch_length(TEST_EPOCH_LENGTH)
-        .sealed_epoch_distance_from_tip(TEST_SEALED_EPOCH_DISTANCE)
-        .build();
-    let mut tf = TestFramework::builder(&mut rng).with_chain_config(chain_config).build();
+    let (vrf_sk, vrf_pk) = VRFPrivateKey::new_from_rng(&mut rng, VRFKeyKind::Schnorrkel);
 
     // create initial chain: genesis <- block_a
-    let (vrf_sk, stake_pool_data) = create_stake_pool_data(&mut rng);
-    let (stake_pool_outpoint, pool_id) =
-        add_block_with_stake_pool(&mut rng, &mut tf, stake_pool_data);
+    let (mut tf, stake_pool_outpoint, pool_id) =
+        setup_test_chain_with_staked_pool(&mut rng, vrf_pk.clone());
     let block_a_id = tf.best_block_id();
 
     // prepare and process block_b with StakePool -> ProduceBlockFromStake kernel
     let initial_randomness = tf.chainstate.get_chain_config().initial_randomness();
-    let pos_data = create_pos_data(
-        &mut tf,
+    let sealed_pool_balance =
+        PoSAccountingStorageRead::<SealedStorageTag>::get_pool_balance(&tf.storage, pool_id)
+            .unwrap()
+            .unwrap();
+    let (pos_data, block_timestamp) = mine(
+        BlockTimestamp::from_duration_since_epoch(tf.current_time()),
         stake_pool_outpoint,
+        &vrf_pk,
         &vrf_sk,
         // no epoch is sealed yet so use initial randomness
-        initial_randomness,
+        PoSRandomness::new(initial_randomness),
         pool_id,
+        sealed_pool_balance,
         1,
-    );
+    )
+    .expect("should be able to mine");
     let reward_output = TxOutput::new(
         OutputValue::Coin(Amount::from_atoms(1)),
         OutputPurpose::ProduceBlockFromStake(anyonecanspend_address(), pool_id),
@@ -910,20 +903,28 @@ fn check_pool_balance_after_reorg(#[case] seed: Seed) {
     tf.make_block_builder()
         .with_consensus_data(ConsensusData::PoS(Box::new(pos_data)))
         .with_reward(vec![reward_output])
+        .with_timestamp(block_timestamp)
         .build_and_process()
         .unwrap();
 
     // prepare and process block_c with ProduceBlockFromStake -> ProduceBlockFromStake kernel
     let block_a_reward_outpoint = OutPoint::new(OutPointSourceId::BlockReward(block_a_id), 0);
-    let pos_data = create_pos_data(
-        &mut tf,
+    let sealed_pool_balance =
+        PoSAccountingStorageRead::<SealedStorageTag>::get_pool_balance(&tf.storage, pool_id)
+            .unwrap()
+            .unwrap();
+    let (pos_data, block_timestamp) = mine(
+        block_timestamp,
         block_a_reward_outpoint,
+        &vrf_pk,
         &vrf_sk,
         // no epoch is sealed yet so use initial randomness
-        initial_randomness,
+        PoSRandomness::new(initial_randomness),
         pool_id,
+        sealed_pool_balance,
         1,
-    );
+    )
+    .expect("should be able to mine");
     let reward_output = TxOutput::new(
         OutputValue::Coin(Amount::from_atoms(1)),
         OutputPurpose::ProduceBlockFromStake(anyonecanspend_address(), pool_id),
@@ -932,6 +933,7 @@ fn check_pool_balance_after_reorg(#[case] seed: Seed) {
         .make_block_builder()
         .with_consensus_data(ConsensusData::PoS(Box::new(pos_data)))
         .with_reward(vec![reward_output])
+        .with_timestamp(block_timestamp)
         .with_parent(block_a_id)
         .build_and_process()
         .unwrap()
@@ -946,14 +948,21 @@ fn check_pool_balance_after_reorg(#[case] seed: Seed) {
     // both sealed epoch and pre block randomness can be used
     let sealed_epoch_randomness =
         tf.storage.transaction_ro().unwrap().get_epoch_data(1).unwrap().unwrap();
-    let pos_data = create_pos_data(
-        &mut tf,
+    let sealed_pool_balance =
+        PoSAccountingStorageRead::<SealedStorageTag>::get_pool_balance(&tf.storage, pool_id)
+            .unwrap()
+            .unwrap();
+    let (pos_data, block_timestamp) = mine(
+        block_timestamp,
         block_3_reward_outpoint,
+        &vrf_pk,
         &vrf_sk,
-        sealed_epoch_randomness.randomness().value(),
+        *sealed_epoch_randomness.randomness(),
         pool_id,
+        sealed_pool_balance,
         2,
-    );
+    )
+    .expect("should be able to mine");
     let reward_output = TxOutput::new(
         OutputValue::Coin(Amount::from_atoms(1)),
         OutputPurpose::ProduceBlockFromStake(anyonecanspend_address(), pool_id),
@@ -961,6 +970,7 @@ fn check_pool_balance_after_reorg(#[case] seed: Seed) {
     tf.make_block_builder()
         .with_consensus_data(ConsensusData::PoS(Box::new(pos_data)))
         .with_reward(vec![reward_output])
+        .with_timestamp(block_timestamp)
         .with_parent((*block_c_index.block_id()).into())
         .build_and_process()
         .unwrap();
diff --git a/chainstate/tx-verifier/src/transaction_verifier/tests/hierarchy_read.rs b/chainstate/tx-verifier/src/transaction_verifier/tests/hierarchy_read.rs
index bb961ea83..72f72adc1 100644
--- a/chainstate/tx-verifier/src/transaction_verifier/tests/hierarchy_read.rs
+++ b/chainstate/tx-verifier/src/transaction_verifier/tests/hierarchy_read.rs
@@ -18,7 +18,6 @@ use std::sync::Arc;
 use crate::transaction_verifier::token_issuance_cache::{CachedAuxDataOp, CachedTokenIndexOp};
 
 use super::*;
-use chainstate_types::GenBlockIndex;
 use common::{
     chain::{
         config::Builder as ConfigBuilder, tokens::TokenAuxiliaryData, TxMainChainIndex,