From 45ea805620b042c546a19d91677fffedbbc01d8a Mon Sep 17 00:00:00 2001
From: Luke Parker <lukeparker5132@gmail.com>
Date: Mon, 21 Aug 2023 01:54:25 -0400
Subject: [PATCH] Use an optimistic (and twice as long in the worst case) sleep
 for KeyGen

Possible due to Substrate having an RPC.
---
 tests/coordinator/src/lib.rs       | 13 +++++++++----
 tests/coordinator/src/tests/mod.rs | 26 ++++++++++++++++++++++----
 2 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/tests/coordinator/src/lib.rs b/tests/coordinator/src/lib.rs
index e09743bc..d46a2135 100644
--- a/tests/coordinator/src/lib.rs
+++ b/tests/coordinator/src/lib.rs
@@ -15,6 +15,8 @@ use serai_client::primitives::NetworkId;
 use messages::{CoordinatorMessage, ProcessorMessage};
 use serai_message_queue::{Service, Metadata, client::MessageQueue};
 
+use serai_client::Serai;
+
 use dockertest::{
   PullPolicy, Image, LogAction, LogPolicy, LogSource, LogOptions, StartPolicy, Composition,
   DockerOperations,
@@ -136,8 +138,7 @@ pub fn coordinator_stack(name: &str) -> (Handles, <Ristretto as Ciphersuite>::F,
 pub struct Processor {
   network: NetworkId,
 
-  #[allow(unused)]
-  serai_handle: String,
+  serai_rpc: String,
   #[allow(unused)]
   message_queue_handle: String,
   #[allow(unused)]
@@ -164,7 +165,7 @@ impl Processor {
     // Bound execution to 60 seconds
     for _ in 0 .. 60 {
       tokio::time::sleep(Duration::from_secs(1)).await;
-      let Ok(client) = serai_client::Serai::new(&serai_rpc).await else { continue };
+      let Ok(client) = Serai::new(&serai_rpc).await else { continue };
       if client.get_latest_block_hash().await.is_err() {
         continue;
       }
@@ -177,7 +178,7 @@ impl Processor {
     Processor {
       network,
 
-      serai_handle: handles.0,
+      serai_rpc,
       message_queue_handle: handles.1,
       coordinator_handle: handles.2,
 
@@ -191,6 +192,10 @@ impl Processor {
     }
   }
 
+  pub async fn serai(&self) -> Serai {
+    Serai::new(&self.serai_rpc).await.unwrap()
+  }
+
   /// Send a message to a processor as its coordinator.
   pub async fn send_message(&mut self, msg: impl Into<ProcessorMessage>) {
     let msg: ProcessorMessage = msg.into();
diff --git a/tests/coordinator/src/tests/mod.rs b/tests/coordinator/src/tests/mod.rs
index b0b26ec5..1c016fda 100644
--- a/tests/coordinator/src/tests/mod.rs
+++ b/tests/coordinator/src/tests/mod.rs
@@ -117,6 +117,9 @@ async fn key_gen_test() {
         processor.send_message(messages::key_gen::ProcessorMessage::Shares { id, shares }).await;
       }
 
+      let serai = processors[0].serai().await;
+      let mut last_serai_block = serai.get_latest_block().await.unwrap().number();
+
       tokio::time::sleep(Duration::from_secs(20)).await;
       if std::env::var("GITHUB_CI") == Ok("true".to_string()) {
         tokio::time::sleep(Duration::from_secs(20)).await;
@@ -149,10 +152,25 @@ async fn key_gen_test() {
       }
 
       // Sleeps for longer since we need to wait for a Substrate block as well
-      // TODO: Replace this with Substrate RPC checks and a much smaller sleep
-      tokio::time::sleep(Duration::from_secs(60)).await;
-      if std::env::var("GITHUB_CI") == Ok("true".to_string()) {
-        tokio::time::sleep(Duration::from_secs(60)).await;
+      'outer: for _ in 0 .. 20 {
+        tokio::time::sleep(Duration::from_secs(6)).await;
+        if std::env::var("GITHUB_CI") == Ok("true".to_string()) {
+          tokio::time::sleep(Duration::from_secs(6)).await;
+        }
+
+        while last_serai_block <= serai.get_latest_block().await.unwrap().number() {
+          if !serai
+            .get_key_gen_events(
+              serai.get_block_by_number(last_serai_block).await.unwrap().unwrap().hash(),
+            )
+            .await
+            .unwrap()
+            .is_empty()
+          {
+            break 'outer;
+          }
+          last_serai_block += 1;
+        }
       }
       let mut message = None;
       for processor in processors.iter_mut() {