From ee888c6567f1fba094970c3a68ebd3582cd4e350 Mon Sep 17 00:00:00 2001
From: Dzejkop <jakubtrad@gmail.com>
Date: Thu, 29 Feb 2024 14:59:47 +0100
Subject: [PATCH] Fix and refactor tests

---
 tests/common/mod.rs        | 23 ++++++++++++++++++++++-
 tests/send_many_txs.rs     | 14 ++------------
 tests/send_too_many_txs.rs | 29 +++++++++++++++++++++++++++--
 3 files changed, 51 insertions(+), 15 deletions(-)

diff --git a/tests/common/mod.rs b/tests/common/mod.rs
index 3d5f121..c514749 100644
--- a/tests/common/mod.rs
+++ b/tests/common/mod.rs
@@ -1,12 +1,13 @@
 #![allow(dead_code)] // Needed because this module is imported as module by many test crates
 
 use std::sync::Arc;
+use std::time::Duration;
 
 use ethers::core::k256::ecdsa::SigningKey;
 use ethers::middleware::SignerMiddleware;
 use ethers::providers::{Http, Middleware, Provider};
 use ethers::signers::{LocalWallet, Signer};
-use ethers::types::{Address, H160};
+use ethers::types::{Address, H160, U256};
 use postgres_docker_utils::DockerContainerGuard;
 use tracing::level_filters::LevelFilter;
 use tracing_subscriber::layer::SubscriberExt;
@@ -104,3 +105,23 @@ pub async fn setup_provider(
 
     Ok(provider)
 }
+
+pub async fn await_balance(
+    provider: &Provider<Http>,
+    expected_balance: U256,
+    address: Address,
+) -> eyre::Result<()> {
+    for _ in 0..50 {
+        let balance = provider.get_balance(address, None).await?;
+
+        tracing::info!(?balance, ?expected_balance, "Checking balance");
+
+        if balance >= expected_balance {
+            return Ok(());
+        } else {
+            tokio::time::sleep(Duration::from_secs(5)).await;
+        }
+    }
+
+    eyre::bail!("Balance did not reach expected value");
+}
diff --git a/tests/send_many_txs.rs b/tests/send_many_txs.rs
index 4325423..bb273f6 100644
--- a/tests/send_many_txs.rs
+++ b/tests/send_many_txs.rs
@@ -47,17 +47,7 @@ async fn send_many_txs() -> eyre::Result<()> {
     }
 
     let expected_balance = value * num_transfers;
-    for _ in 0..50 {
-        let balance = provider.get_balance(ARBITRARY_ADDRESS, None).await?;
+    await_balance(&provider, expected_balance, ARBITRARY_ADDRESS).await?;
 
-        tracing::info!(?balance, ?expected_balance, "Checking balance");
-
-        if balance == expected_balance {
-            return Ok(());
-        } else {
-            tokio::time::sleep(Duration::from_secs(5)).await;
-        }
-    }
-
-    panic!("Transactions were not sent")
+    Ok(())
 }
diff --git a/tests/send_too_many_txs.rs b/tests/send_too_many_txs.rs
index 35b0959..226cc43 100644
--- a/tests/send_too_many_txs.rs
+++ b/tests/send_too_many_txs.rs
@@ -18,12 +18,12 @@ async fn send_too_many_txs() -> eyre::Result<()> {
     let (_service, client) =
         ServiceBuilder::default().build(&anvil, &db_url).await?;
 
-    let CreateApiKeyResponse { api_key: _api_key } =
+    let CreateApiKeyResponse { api_key } =
         client.create_relayer_api_key(DEFAULT_RELAYER_ID).await?;
 
     let CreateRelayerResponse {
         relayer_id: secondary_relayer_id,
-        address: _secondary_relayer_address,
+        address: secondary_relayer_address,
     } = client
         .create_relayer(&CreateRelayerRequest {
             name: "Secondary Relayer".to_string(),
@@ -88,5 +88,30 @@ async fn send_too_many_txs() -> eyre::Result<()> {
         result
     );
 
+    // Accumulate total value + gas budget
+    let send_value = value * (MAX_QUEUED_TXS + 1);
+    let total_required_value = send_value + parse_units("1", "ether")?;
+
+    client
+        .send_tx(
+            &api_key,
+            &SendTxRequest {
+                to: secondary_relayer_address,
+                value: total_required_value,
+                data: None,
+                gas_limit: U256::from(21_000),
+                priority: TransactionPriority::Regular,
+                tx_id: None,
+            },
+        )
+        .await?;
+
+    tracing::info!("Waiting for secondary relayer balance");
+    await_balance(&provider, total_required_value, secondary_relayer_address)
+        .await?;
+
+    tracing::info!("Waiting for queued up txs to be processed");
+    await_balance(&provider, send_value, ARBITRARY_ADDRESS).await?;
+
     Ok(())
 }