From e0dc5d29ada6e3cf290c79595965bbcd1f67db3c Mon Sep 17 00:00:00 2001 From: Luke Parker Date: Sun, 23 Apr 2023 01:52:19 -0400 Subject: [PATCH] Tributary test wait_for_tx_inclusion function --- coordinator/src/tests/tributary/chain.rs | 40 +++++++++++++++++++++++- coordinator/src/tests/tributary/dkg.rs | 22 +++++-------- coordinator/src/tests/tributary/tx.rs | 28 +++++------------ 3 files changed, 53 insertions(+), 37 deletions(-) diff --git a/coordinator/src/tests/tributary/chain.rs b/coordinator/src/tests/tributary/chain.rs index 416e9329..ec6cb69f 100644 --- a/coordinator/src/tests/tributary/chain.rs +++ b/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, + 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); diff --git a/coordinator/src/tests/tributary/dkg.rs b/coordinator/src/tests/tributary/dkg.rs index 702cf79c..a618d1de 100644 --- a/coordinator/src/tests/tributary/dkg.rs +++ b/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::::block_time()).into(), )) diff --git a/coordinator/src/tests/tributary/tx.rs b/coordinator/src/tests/tributary/tx.rs index 973e98b2..9c6ea156 100644 --- a/coordinator/src/tests/tributary/tx.rs +++ b/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::::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::::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()]); } }