From f56827500181cadd24abb4f7757cdadc5444843c Mon Sep 17 00:00:00 2001
From: Aaron Blankstein <aaron@hiro.so>
Date: Fri, 17 Jan 2025 15:15:30 -0600
Subject: [PATCH] fix flakiness in forked-tenure-is-ignored

---
 .../src/tests/nakamoto_integrations.rs        | 45 ++++++++++++++-----
 1 file changed, 33 insertions(+), 12 deletions(-)

diff --git a/testnet/stacks-node/src/tests/nakamoto_integrations.rs b/testnet/stacks-node/src/tests/nakamoto_integrations.rs
index 37b3444902..dbf7f64017 100644
--- a/testnet/stacks-node/src/tests/nakamoto_integrations.rs
+++ b/testnet/stacks-node/src/tests/nakamoto_integrations.rs
@@ -47,7 +47,7 @@ use stacks::chainstate::stacks::address::{PoxAddress, StacksAddressExtensions};
 use stacks::chainstate::stacks::boot::{
     MINERS_NAME, SIGNERS_VOTING_FUNCTION_NAME, SIGNERS_VOTING_NAME,
 };
-use stacks::chainstate::stacks::db::StacksChainState;
+use stacks::chainstate::stacks::db::{StacksChainState, StacksHeaderInfo};
 use stacks::chainstate::stacks::miner::{
     BlockBuilder, BlockLimitFunction, TransactionEvent, TransactionResult, TransactionSuccessEvent,
 };
@@ -304,6 +304,30 @@ pub fn get_stackerdb_slot_version(
     })
 }
 
+pub fn get_last_block_in_current_tenure(
+    sortdb: &SortitionDB,
+    chainstate: &StacksChainState,
+) -> Option<StacksHeaderInfo> {
+    let ch = SortitionDB::get_canonical_burn_chain_tip(sortdb.conn())
+        .unwrap()
+        .consensus_hash;
+    let mut tenure_blocks = test_observer::get_blocks();
+    tenure_blocks.retain(|block| {
+        let consensus_hash = block.get("consensus_hash").unwrap().as_str().unwrap();
+        consensus_hash == format!("0x{ch}")
+    });
+    let last_block = tenure_blocks.last()?.clone();
+    let last_block_id = StacksBlockId::from_hex(
+        &last_block
+            .get("index_block_hash")
+            .unwrap()
+            .as_str()
+            .unwrap()[2..],
+    )
+    .unwrap();
+    NakamotoChainState::get_block_header(chainstate.db(), &last_block_id).unwrap()
+}
+
 pub fn add_initial_balances(
     conf: &mut Config,
     accounts: usize,
@@ -5023,7 +5047,7 @@ fn forked_tenure_is_ignored() {
         thread::sleep(Duration::from_secs(1));
     }
 
-    info!("Tenure B broadcasted but did not process a block. Issue the next bitcon block and unstall block commits.");
+    info!("Tenure B broadcasted but did not process a block. Issue the next bitcoin block and unstall block commits.");
 
     // the block will be stored, not processed, so load it out of staging
     let tip_sn = SortitionDB::get_canonical_burn_chain_tip(sortdb.conn())
@@ -5070,16 +5094,17 @@ fn forked_tenure_is_ignored() {
             .lock()
             .expect("Mutex poisoned")
             .get_stacks_blocks_processed();
+        let block_in_tenure = get_last_block_in_current_tenure(&sortdb, &chainstate).is_some();
         Ok(commits_count > commits_before
             && blocks_count > blocks_before
-            && blocks_processed > blocks_processed_before)
+            && blocks_processed > blocks_processed_before
+            && block_in_tenure)
     })
     .unwrap();
 
     info!("Tenure C produced a block!");
-    let block_tenure_c = NakamotoChainState::get_canonical_block_header(chainstate.db(), &sortdb)
-        .unwrap()
-        .unwrap();
+
+    let block_tenure_c = get_last_block_in_current_tenure(&sortdb, &chainstate).unwrap();
     let blocks = test_observer::get_mined_nakamoto_blocks();
     let block_c = blocks.last().unwrap();
     info!("Tenure C tip block: {}", &block_tenure_c.index_block_hash());
@@ -5132,9 +5157,7 @@ fn forked_tenure_is_ignored() {
 
     info!("Tenure C produced a second block!");
 
-    let block_2_tenure_c = NakamotoChainState::get_canonical_block_header(chainstate.db(), &sortdb)
-        .unwrap()
-        .unwrap();
+    let block_2_tenure_c = get_last_block_in_current_tenure(&sortdb, &chainstate).unwrap();
     let blocks = test_observer::get_mined_nakamoto_blocks();
     let block_2_c = blocks.last().unwrap();
 
@@ -5165,9 +5188,7 @@ fn forked_tenure_is_ignored() {
     })
     .unwrap();
 
-    let block_tenure_d = NakamotoChainState::get_canonical_block_header(chainstate.db(), &sortdb)
-        .unwrap()
-        .unwrap();
+    let block_tenure_d = get_last_block_in_current_tenure(&sortdb, &chainstate).unwrap();
     let blocks = test_observer::get_mined_nakamoto_blocks();
     let block_d = blocks.last().unwrap();