Tributary test wait_for_tx_inclusion function

coordinator/src/tests/tributary/chain.rs

@@ -20,7 +20,7 @@ use tokio::time::sleep;
 
 use serai_db::MemDb;
 
-use tributary::Tributary;
+use tributary::{Transaction as TransactionTrait, Tributary};
 
 use crate::{
   P2pMessageKind, P2p, LocalP2p,
@@ -108,6 +108,44 @@ pub async fn run_tributaries(
   }
 }
 
+pub async fn wait_for_tx_inclusion(
+  tributary: &Tributary<MemDb, Transaction, LocalP2p>,
+  mut last_checked: [u8; 32],
+  hash: [u8; 32],
+) -> [u8; 32] {
+  loop {
+    let tip = tributary.tip();
+    if tip == last_checked {
+      sleep(Duration::from_secs(1)).await;
+      continue;
+    }
+
+    let mut queue = vec![tributary.block(&tip).unwrap()];
+    let mut block = None;
+    while {
+      let parent = queue.last().unwrap().parent();
+      if parent == tributary.genesis() {
+        false
+      } else {
+        block = Some(tributary.block(&parent).unwrap());
+        block.as_ref().unwrap().hash() != last_checked
+      }
+    } {
+      queue.push(block.take().unwrap());
+    }
+
+    while let Some(block) = queue.pop() {
+      for tx in &block.transactions {
+        if tx.hash() == hash {
+          return block.hash();
+        }
+      }
+    }
+
+    last_checked = tip;
+  }
+}
+
 #[tokio::test]
 async fn tributary_test() {
   let keys = new_keys(&mut OsRng);

coordinator/src/tests/tributary/dkg.rs

@@ -15,13 +15,13 @@ use processor_messages::{
   CoordinatorMessage,
 };
 
-use tributary::Tributary;
+use tributary::{Transaction as TransactionTrait, Tributary};
 
 use crate::{
   processor::MemProcessor,
   LocalP2p,
   tributary::{TributaryDb, Transaction, TributarySpec, scanner::handle_new_blocks},
-  tests::tributary::{new_keys, new_spec, new_tributaries, run_tributaries},
+  tests::tributary::{new_keys, new_spec, new_tributaries, run_tributaries, wait_for_tx_inclusion},
 };
 
 #[tokio::test]
@@ -47,25 +47,16 @@ async fn dkg_commitments_test() {
     txs.push(tx);
   }
 
-  let mut last_block = tributaries[0].1.tip();
+  let block_before_tx = tributaries[0].1.tip();
 
   // Publish all commitments but one
   for (i, tx) in txs.iter().enumerate().skip(1) {
     assert!(tributaries[i].1.add_transaction(tx.clone()).await);
   }
 
-  // Wait until these were included
-  let mut included = 0;
-  while included != (txs.len() - 1) {
-    let tributary = &tributaries[0].1;
-    let tip = tributary.tip();
-    if tip == last_block {
-      sleep(Duration::from_secs(1)).await;
-      continue;
-    }
-    last_block = tip;
-
-    included += tributary.block(&last_block).unwrap().transactions.len();
+  // Wait until these are included
+  for tx in txs.iter().skip(1) {
+    wait_for_tx_inclusion(&tributaries[0].1, block_before_tx, tx.hash()).await;
   }
 
   let expected_msg = CoordinatorMessage::KeyGen(key_gen::CoordinatorMessage::Commitments {
@@ -104,6 +95,7 @@ async fn dkg_commitments_test() {
 
     // Publish the last commitment
     assert!(tributaries[0].1.add_transaction(txs[0].clone()).await);
+    wait_for_tx_inclusion(&tributaries[0].1, last_block, txs[0].hash()).await;
     sleep(Duration::from_secs(
       (2 * Tributary::<MemDb, Transaction, LocalP2p>::block_time()).into(),
     ))

coordinator/src/tests/tributary/tx.rs

@@ -6,12 +6,12 @@ use tokio::time::sleep;
 
 use serai_db::MemDb;
 
-use tributary::Tributary;
+use tributary::{Transaction as TransactionTrait, Tributary};
 
 use crate::{
   LocalP2p,
   tributary::Transaction,
-  tests::tributary::{new_keys, new_spec, new_tributaries, run_tributaries},
+  tests::tributary::{new_keys, new_spec, new_tributaries, run_tributaries, wait_for_tx_inclusion},
 };
 
 #[tokio::test]
@@ -34,33 +34,19 @@ async fn tx_test() {
   OsRng.fill_bytes(&mut commitments);
 
   // Create the TX with a null signature so we can get its sig hash
+  let block_before_tx = tributaries[sender].1.tip();
   let mut tx =
     Transaction::DkgCommitments(attempt, commitments.clone(), Transaction::empty_signed());
   tx.sign(&mut OsRng, spec.genesis(), &key, 0);
 
   assert!(tributaries[sender].1.add_transaction(tx.clone()).await);
-  // Sleep for two blocks
-  sleep(Duration::from_secs((2 * Tributary::<MemDb, Transaction, LocalP2p>::block_time()).into()))
-    .await;
+  let included_in = wait_for_tx_inclusion(&tributaries[sender].1, block_before_tx, tx.hash()).await;
+  // Also sleep for the block time to ensure the block is synced around before we run checks on it
+  sleep(Duration::from_secs(Tributary::<MemDb, Transaction, LocalP2p>::block_time().into())).await;
 
   // All tributaries should have acknowledged this transaction in a block
-  let mut included_in = None;
   for (_, tributary) in tributaries {
-    if included_in.is_none() {
-      let mut found = tributary.tip();
-
-      let mut block;
-      while {
-        block = tributary.block(&found).unwrap();
-        block.transactions.is_empty()
-      } {
-        found = block.parent();
-      }
-
-      included_in = Some(found);
-    }
-
-    let block = tributary.block(&included_in.unwrap()).unwrap();
+    let block = tributary.block(&included_in).unwrap();
     assert_eq!(block.transactions, vec![tx.clone()]);
   }
 }