From 91b3d0c83adc0f853375522728e48c29aa9e2b84 Mon Sep 17 00:00:00 2001
From: Jesse Abramowitz <jesse@entropy.xyz>
Date: Mon, 28 Oct 2024 10:07:16 -0400
Subject: [PATCH] add run to all blocks

---
 .../tests/register_sign_reshare_sign.rs       | 27 ++++++++++++++-----
 1 file changed, 20 insertions(+), 7 deletions(-)

diff --git a/crates/threshold-signature-server/tests/register_sign_reshare_sign.rs b/crates/threshold-signature-server/tests/register_sign_reshare_sign.rs
index cad941af9..59ef99e24 100644
--- a/crates/threshold-signature-server/tests/register_sign_reshare_sign.rs
+++ b/crates/threshold-signature-server/tests/register_sign_reshare_sign.rs
@@ -120,7 +120,7 @@ async fn integration_test_register_sign_reshare_sign() {
     );
 
     // Do a reshare
-    do_reshare(&api, &rpc).await;
+    do_reshare(&api, &rpc, &other_rpcs).await;
 
     // Sign a message again
     let recoverable_signature = test_client::sign(
@@ -148,7 +148,11 @@ async fn integration_test_register_sign_reshare_sign() {
     );
 }
 
-async fn do_reshare(api: &OnlineClient<EntropyConfig>, rpc: &LegacyRpcMethods<EntropyConfig>) {
+async fn do_reshare(
+    api: &OnlineClient<EntropyConfig>,
+    rpc: &LegacyRpcMethods<EntropyConfig>,
+    rpcs: &[LegacyRpcMethods<EntropyConfig>],
+) {
     // Get current signers
     let signer_query = entropy::storage().staking_extension().signers();
     let signer_stash_accounts = query_chain(&api, &rpc, signer_query, None).await.unwrap().unwrap();
@@ -167,13 +171,12 @@ async fn do_reshare(api: &OnlineClient<EntropyConfig>, rpc: &LegacyRpcMethods<En
         new_signers: reshare_data.new_signers.into_iter().map(|s| s.to_vec()).collect(),
         block_number: block_number - 1,
     };
-
-    run_to_block(&rpc, block_number).await;
+    let ips = vec![3002, 3003, 3004];
+    run_to_all_blocks(rpcs, block_number).await;
     // Send the OCW message to all TS servers who don't have a chain node
     let client = reqwest::Client::new();
     let response_results = join_all(
-        [3002, 3003, 3004]
-            .iter()
+        ips.iter()
             .map(|port| {
                 client
                     .post(format!("http://127.0.0.1:{}/validator/reshare", port))
@@ -183,7 +186,8 @@ async fn do_reshare(api: &OnlineClient<EntropyConfig>, rpc: &LegacyRpcMethods<En
             .collect::<Vec<_>>(),
     )
     .await;
-    for response_result in response_results {
+    for (i, response_result) in response_results.into_iter().enumerate() {
+        dbg!(ips[i]);
         assert_eq!(response_result.unwrap().text().await.unwrap(), "");
     }
 
@@ -221,3 +225,12 @@ async fn do_reshare(api: &OnlineClient<EntropyConfig>, rpc: &LegacyRpcMethods<En
     let new: HashSet<[u8; 32]> = new_signer_stash_accounts.iter().map(|s| s.0).collect();
     assert_ne!(old, new);
 }
+
+pub async fn run_to_all_blocks(rpcs: &[LegacyRpcMethods<EntropyConfig>], block_run: u32) {
+    let mut current_block = 0;
+    for rpc in rpcs {
+        while current_block < block_run {
+            current_block = rpc.chain_get_header(None).await.unwrap().unwrap().number;
+        }
+    }
+}