Merge pull request #1128 from mintlayer/fix/use_pos_block_builder_more

Rewrite more unit tests with PoSBlockBuilder

chainstate/test-framework/src/lib.rs

@@ -31,8 +31,9 @@ pub type TestChainstate = Box<dyn chainstate::chainstate_interface::ChainstateIn
 
 pub use {
     crate::utils::{
-        anyonecanspend_address, create_chain_config_with_staking_pool, empty_witness,
-        get_output_value, get_target_block_time, pos_mine, produce_kernel_signature,
+        anyonecanspend_address, create_chain_config_with_default_staking_pool,
+        create_chain_config_with_staking_pool, empty_witness, get_output_value,
+        get_target_block_time, pos_mine, produce_kernel_signature,
     },
     block_builder::BlockBuilder,
     framework::TestFramework,

chainstate/test-framework/src/pos_block_builder.rs

@@ -32,7 +32,7 @@ use common::{
         signed_transaction::SignedTransaction,
         Block, Destination, GenBlock, PoolId, RequiredConsensus, TxOutput, UtxoOutPoint,
     },
-    primitives::{Id, Idable, H256},
+    primitives::{Amount, Id, Idable, H256},
 };
 use crypto::{
     key::{KeyKind, PrivateKey, PublicKey},
@@ -55,6 +55,9 @@ pub struct PoSBlockBuilder<'f> {
 
     staker_sk: PrivateKey,
     staker_vrf_sk: VRFPrivateKey,
+
+    randomness: Option<PoSRandomness>,
+    stake_pool_balance: Option<Amount>,
 }
 
 impl<'f> PoSBlockBuilder<'f> {
@@ -91,6 +94,8 @@ impl<'f> PoSBlockBuilder<'f> {
             kernel_input_outpoint: None,
             staker_sk,
             staker_vrf_sk,
+            randomness: None,
+            stake_pool_balance: None,
         }
     }
 
@@ -144,6 +149,16 @@ impl<'f> PoSBlockBuilder<'f> {
         self
     }
 
+    pub fn with_randomness(mut self, randomness: PoSRandomness) -> Self {
+        self.randomness = Some(randomness);
+        self
+    }
+
+    pub fn with_stake_pool_balance(mut self, balance: Amount) -> Self {
+        self.stake_pool_balance = Some(balance);
+        self
+    }
+
     pub fn with_kernel_input(mut self, outpoint: UtxoOutPoint) -> Self {
         self.kernel_input_outpoint = Some(outpoint);
         self
@@ -272,14 +287,19 @@ impl<'f> PoSBlockBuilder<'f> {
             }
         };
         let chain_config = self.framework.chainstate.get_chain_config().as_ref();
-        let randomness = match chain_config.sealed_epoch_index(&new_block_height) {
-            Some(epoch) => {
-                *self.framework.storage.get_epoch_data(epoch).unwrap().unwrap().randomness()
+        let epoch_index = chain_config.epoch_index_from_height(&new_block_height);
+
+        let randomness = self.randomness.unwrap_or_else(|| {
+            match chain_config.sealed_epoch_index(&new_block_height) {
+                Some(epoch) => {
+                    *self.framework.storage.get_epoch_data(epoch).unwrap().unwrap().randomness()
+                }
+                None => PoSRandomness::new(chain_config.initial_randomness()),
             }
-            None => PoSRandomness::new(chain_config.initial_randomness()),
-        };
-        let pool_balance =
-            self.framework.chainstate.get_stake_pool_balance(staking_pool).unwrap().unwrap();
+        });
+        let pool_balance = self.stake_pool_balance.unwrap_or_else(|| {
+            self.framework.chainstate.get_stake_pool_balance(staking_pool).unwrap().unwrap()
+        });
 
         pos_mine(
             BlockTimestamp::from_duration_since_epoch(self.framework.current_time()),
@@ -289,7 +309,7 @@ impl<'f> PoSBlockBuilder<'f> {
             randomness,
             staking_pool,
             pool_balance,
-            0,
+            epoch_index,
             current_difficulty.into(),
         )
         .unwrap()

chainstate/test-framework/src/utils.rs

@@ -21,7 +21,7 @@ use chainstate_types::pos_randomness::PoSRandomness;
 use common::{
     chain::{
         block::{consensus_data::PoSData, timestamp::BlockTimestamp, BlockRewardTransactable},
-        config::{Builder as ConfigBuilder, ChainType, EpochIndex},
+        config::{create_unit_test_config, Builder as ConfigBuilder, ChainType, EpochIndex},
         create_unittest_pos_config,
         output_value::OutputValue,
         signature::{
@@ -34,11 +34,11 @@ use common::{
         OutPointSourceId, PoolId, RequiredConsensus, TxInput, TxOutput, UpgradeVersion,
         UtxoOutPoint,
     },
-    primitives::{Amount, BlockHeight, Compact, Id, Idable},
+    primitives::{per_thousand::PerThousand, Amount, BlockHeight, Compact, Id, Idable, H256},
     Uint256,
 };
 use crypto::{
-    key::PrivateKey,
+    key::{PrivateKey, PublicKey},
     random::{CryptoRng, Rng},
     vrf::{VRFPrivateKey, VRFPublicKey},
 };
@@ -317,6 +317,33 @@ pub fn get_target_block_time(chain_config: &ChainConfig, block_height: BlockHeig
     }
 }
 
+pub fn create_chain_config_with_default_staking_pool(
+    rng: &mut impl Rng,
+    staking_pk: PublicKey,
+    vrf_pk: VRFPublicKey,
+) -> (ConfigBuilder, PoolId) {
+    let stake_amount = create_unit_test_config().min_stake_pool_pledge();
+    let mint_amount = Amount::from_atoms(100_000_000 * common::chain::Mlt::ATOMS_PER_MLT);
+
+    let genesis_pool_id = PoolId::new(H256::random_using(rng));
+    let genesis_stake_pool_data = StakePoolData::new(
+        stake_amount,
+        Destination::PublicKey(staking_pk.clone()),
+        vrf_pk,
+        Destination::PublicKey(staking_pk),
+        PerThousand::new(1000).unwrap(),
+        Amount::ZERO,
+    );
+
+    let chain_config = create_chain_config_with_staking_pool(
+        mint_amount,
+        genesis_pool_id,
+        genesis_stake_pool_data,
+    );
+
+    (chain_config, genesis_pool_id)
+}
+
 pub fn create_chain_config_with_staking_pool(
     mint_amount: Amount,
     pool_id: PoolId,

chainstate/test-suite/src/tests/pos_accounting_reorg.rs

@@ -481,3 +481,73 @@ fn in_memory_reorg_disconnect_create_pool(#[case] seed: Seed) {
         tf.best_block_id()
     );
 }
+
+// Use 2 chainstates to produce 2 branches from a common genesis using PoS.
+// Chainstate1: genesis <- a
+// Chainstate2: genesis <- b <- c
+// Process blocks from Chainstate2 using Chainstate1 and check that reorg happens.
+#[rstest]
+#[trace]
+#[case(Seed::from_entropy())]
+fn pos_reorg_simple(#[case] seed: Seed) {
+    let mut rng = make_seedable_rng(seed);
+
+    let (vrf_sk, vrf_pk) = VRFPrivateKey::new_from_rng(&mut rng, VRFKeyKind::Schnorrkel);
+    let (staker_sk, staker_pk) = PrivateKey::new_from_rng(&mut rng, KeyKind::Secp256k1Schnorr);
+
+    let (chain_config_builder, _) =
+        chainstate_test_framework::create_chain_config_with_default_staking_pool(
+            &mut rng, staker_pk, vrf_pk,
+        );
+    let chain_config = chain_config_builder.build();
+
+    let target_block_time =
+        chainstate_test_framework::get_target_block_time(&chain_config, BlockHeight::new(1));
+
+    let mut tf1 = TestFramework::builder(&mut rng).with_chain_config(chain_config.clone()).build();
+    let mut tf2 = TestFramework::builder(&mut rng).with_chain_config(chain_config).build();
+
+    tf1.progress_time_seconds_since_epoch(target_block_time.get());
+    tf2.progress_time_seconds_since_epoch(target_block_time.get());
+
+    // Block A
+    tf1.make_pos_block_builder(&mut rng)
+        .with_block_signing_key(staker_sk.clone())
+        .with_stake_spending_key(staker_sk.clone())
+        .with_vrf_key(vrf_sk.clone())
+        .build_and_process()
+        .unwrap();
+
+    // Block B
+    let block_b = tf2
+        .make_pos_block_builder(&mut rng)
+        .with_block_signing_key(staker_sk.clone())
+        .with_stake_spending_key(staker_sk.clone())
+        .with_vrf_key(vrf_sk.clone())
+        .build();
+    tf2.process_block(block_b.clone(), BlockSource::Local).unwrap();
+
+    // Block C
+    let block_c = tf2
+        .make_pos_block_builder(&mut rng)
+        .with_block_signing_key(staker_sk.clone())
+        .with_stake_spending_key(staker_sk)
+        .with_vrf_key(vrf_sk)
+        .build();
+    let block_c_id = block_c.get_id();
+    tf2.process_block(block_c.clone(), BlockSource::Local).unwrap();
+
+    assert_eq!(<Id<GenBlock>>::from(block_c_id), tf2.best_block_id());
+
+    // Try to switch to a new branch
+
+    tf1.chainstate.preliminary_header_check(block_b.header().clone()).unwrap();
+    let block_b = tf1.chainstate.preliminary_block_check(block_b).unwrap();
+    tf1.process_block(block_b, BlockSource::Peer).unwrap();
+
+    tf1.chainstate.preliminary_header_check(block_c.header().clone()).unwrap();
+    let block_c = tf1.chainstate.preliminary_block_check(block_c).unwrap();
+    tf1.process_block(block_c, BlockSource::Peer).unwrap().unwrap();
+
+    assert_eq!(<Id<GenBlock>>::from(block_c_id), tf1.best_block_id());
+}