Add a test for Tributary

Further fleshes out the Tributary testing code.

coordinator/src/tests/tributary/chain.rs

@@ -2,7 +2,7 @@ use std::time::{Duration, SystemTime};
 
 use zeroize::Zeroizing;
 
-use rand_core::{RngCore, OsRng};
+use rand_core::{RngCore, CryptoRng, OsRng};
 
 use ciphersuite::{
   group::{ff::Field, GroupEncoding},
@@ -22,11 +22,27 @@ use serai_db::MemDb;
 
 use tributary::Tributary;
 
-use crate::{P2pMessageKind, P2p, LocalP2p, processor::MemProcessor, tributary::TributarySpec};
+use crate::{
+  P2pMessageKind, P2p, LocalP2p,
+  tributary::{Transaction, TributarySpec},
+};
 
-fn new_spec(keys: &[Zeroizing<<Ristretto as Ciphersuite>::F>]) -> TributarySpec {
+pub fn new_keys<R: RngCore + CryptoRng>(
+  rng: &mut R,
+) -> Vec<Zeroizing<<Ristretto as Ciphersuite>::F>> {
+  let mut keys = vec![];
+  for _ in 0 .. 5 {
+    keys.push(Zeroizing::new(<Ristretto as Ciphersuite>::F::random(&mut *rng)));
+  }
+  keys
+}
+
+pub fn new_spec<R: RngCore + CryptoRng>(
+  rng: &mut R,
+  keys: &[Zeroizing<<Ristretto as Ciphersuite>::F>],
+) -> TributarySpec {
   let mut serai_block = [0; 32];
-  OsRng.fill_bytes(&mut serai_block);
+  rng.fill_bytes(&mut serai_block);
 
   let start_time = SystemTime::now().duration_since(SystemTime::UNIX_EPOCH).unwrap().as_secs();
 
@@ -48,24 +64,16 @@ fn new_spec(keys: &[Zeroizing<<Ristretto as Ciphersuite>::F>]) -> TributarySpec
   TributarySpec::new(serai_block, start_time, set, set_data)
 }
 
-#[tokio::test]
-async fn tributary_test() {
-  let mut keys = vec![];
-  for _ in 0 .. 5 {
-    keys.push(Zeroizing::new(<Ristretto as Ciphersuite>::F::random(&mut OsRng)));
-  }
-
-  let processor = MemProcessor::new();
-
-  let spec = new_spec(&keys);
-
+pub async fn new_tributaries(
+  keys: &[Zeroizing<<Ristretto as Ciphersuite>::F>],
+  spec: &TributarySpec,
+) -> Vec<(LocalP2p, Tributary<MemDb, Transaction, LocalP2p>)> {
   let p2p = LocalP2p::new(keys.len());
-
-  let mut tributaries = vec![];
-
+  let mut res = vec![];
   for (i, key) in keys.iter().enumerate() {
-    tributaries.push(
-      Tributary::<_, crate::tributary::Transaction, _>::new(
+    res.push((
+      p2p[i].clone(),
+      Tributary::<_, Transaction, _>::new(
         MemDb::new(),
         spec.genesis(),
         spec.start_time(),
@@ -75,25 +83,56 @@ async fn tributary_test() {
       )
       .await
       .unwrap(),
-    );
+    ));
   }
+  res
+}
 
-  let mut blocks = 0;
-  let mut last_block = spec.genesis();
-
-  let timeout = SystemTime::now() + Duration::from_secs(70);
-  while (blocks < 10) && (SystemTime::now().duration_since(timeout).is_err()) {
-    for (i, p2p) in p2p.iter().enumerate() {
+pub async fn run_tributaries(
+  mut tributaries: Vec<(LocalP2p, Tributary<MemDb, Transaction, LocalP2p>)>,
+) {
+  loop {
+    for (p2p, tributary) in tributaries.iter_mut() {
       while let Some(msg) = p2p.receive().await {
         match msg.0 {
           P2pMessageKind::Tributary => {
-            tributaries[i].handle_message(&msg.1).await;
+            if tributary.handle_message(&msg.1).await {
+              p2p.broadcast(msg.0, msg.1).await;
+            }
           }
         }
       }
     }
 
-    let tip = tributaries[0].tip();
+    sleep(Duration::from_millis(100)).await;
+  }
+}
+
+#[tokio::test]
+async fn tributary_test() {
+  let keys = new_keys(&mut OsRng);
+  let spec = new_spec(&mut OsRng, &keys);
+
+  let mut tributaries = new_tributaries(&keys, &spec).await;
+
+  let mut blocks = 0;
+  let mut last_block = spec.genesis();
+
+  // Doesn't use run_tributaries as we want to wind these down at a certain point
+  // run_tributaries will run them ad infinitum
+  let timeout = SystemTime::now() + Duration::from_secs(65);
+  while (blocks < 10) && (SystemTime::now().duration_since(timeout).is_err()) {
+    for (p2p, tributary) in tributaries.iter_mut() {
+      while let Some(msg) = p2p.receive().await {
+        match msg.0 {
+          P2pMessageKind::Tributary => {
+            tributary.handle_message(&msg.1).await;
+          }
+        }
+      }
+    }
+
+    let tip = tributaries[0].1.tip();
     if tip != last_block {
       last_block = tip;
       blocks += 1;
@@ -107,11 +146,11 @@ async fn tributary_test() {
   }
 
   // Handle all existing messages
-  for (i, p2p) in p2p.iter().enumerate() {
+  for (p2p, tributary) in tributaries.iter_mut() {
     while let Some(msg) = p2p.receive().await {
       match msg.0 {
         P2pMessageKind::Tributary => {
-          tributaries[i].handle_message(&msg.1).await;
+          tributary.handle_message(&msg.1).await;
         }
       }
     }
@@ -119,7 +158,7 @@ async fn tributary_test() {
 
   // All tributaries should agree on the tip
   let mut final_block = None;
-  for tributary in tributaries {
+  for (_, tributary) in tributaries {
     final_block = final_block.or_else(|| Some(tributary.tip()));
     if tributary.tip() != final_block.unwrap() {
       panic!("tributary had different tip");