feat: staking-cli: concurrent txns, better output

sequencer/src/api.rs

@@ -980,7 +980,7 @@ pub mod test_helpers {
     use itertools::izip;
     use jf_merkle_tree_compat::{MerkleCommitment, MerkleTreeScheme};
     use portpicker::pick_unused_port;
-    use staking_cli::demo::{setup_stake_table_contract_for_test, DelegationConfig};
+    use staking_cli::demo::{DelegationConfig, StakingTransactions};
     use surf_disco::Client;
     use tempfile::TempDir;
     use tide_disco::{error::ServerError, Api, App, Error, StatusCode};
@@ -1217,15 +1217,18 @@ pub mod test_helpers {
             let stake_table_address = contracts
                 .address(Contract::StakeTableProxy)
                 .expect("StakeTableProxy address not found");
-            setup_stake_table_contract_for_test(
+            StakingTransactions::create(
                 l1_url.clone(),
                 &deployer,
                 stake_table_address,
                 network_config.staking_priv_keys(),
                 delegation_config,
             )
             .await
-            .expect("stake table setup failed");
+            .expect("stake table setup failed")
+            .apply_all()
+            .await
+            .expect("send all txns failed");
 
             // enable interval mining with a 1s interval.
             // This ensures that blocks are finalized every second, even when there are no transactions.
@@ -3544,6 +3547,7 @@ mod test {
         let mut target_bh = 0;
         while let Some(header) = headers.next().await {
             let header = header.unwrap();
+            println!("got header with height {}", header.height());
             if header.height() == 0 {
                 continue;
             }
@@ -5825,7 +5829,10 @@ mod test {
             .try_into()?;
             commissions.push((validator, commission, new_commission));
             tracing::info!(%validator, %commission, %new_commission, "Update commission");
-            update_commission(provider, st_addr, new_commission).await?;
+            update_commission(provider, st_addr, new_commission)
+                .await?
+                .get_receipt()
+                .await?;
         }
 
         // wait until new stake table takes effect

sequencer/src/bin/espresso-dev-node.rs

@@ -51,7 +51,7 @@ use sequencer::{
 };
 use sequencer_utils::logging;
 use serde::{Deserialize, Serialize};
-use staking_cli::demo::{setup_stake_table_contract_for_test, DelegationConfig};
+use staking_cli::demo::{DelegationConfig, StakingTransactions};
 use tempfile::NamedTempFile;
 use tide_disco::{error::ServerError, method::ReadState, Api, Error, StatusCode};
 use tokio::spawn;
@@ -536,7 +536,7 @@ async fn main() -> anyhow::Result<()> {
             }
 
             let staking_priv_keys = network_config.staking_priv_keys();
-            setup_stake_table_contract_for_test(
+            StakingTransactions::create(
                 l1_url.clone(),
                 &provider,
                 l1_contracts
@@ -545,6 +545,8 @@ async fn main() -> anyhow::Result<()> {
                 staking_priv_keys,
                 DelegationConfig::default(),
             )
+            .await?
+            .apply_all()
             .await?;
         }
     }

sequencer/src/lib.rs

@@ -695,7 +695,7 @@ pub mod testing {
     use portpicker::pick_unused_port;
     use rand::SeedableRng as _;
     use rand_chacha::ChaCha20Rng;
-    use staking_cli::demo::{setup_stake_table_contract_for_test, DelegationConfig};
+    use staking_cli::demo::{DelegationConfig, StakingTransactions};
     use tokio::spawn;
     use vbs::version::Version;
 
@@ -938,15 +938,18 @@ pub mod testing {
                     let st_addr = contracts
                         .address(Contract::StakeTableProxy)
                         .expect("StakeTableProxy address not found");
-                    setup_stake_table_contract_for_test(
+                    StakingTransactions::create(
                         self.l1_url.clone(),
                         &deployer,
                         st_addr,
                         validators,
                         DelegationConfig::default(),
                     )
                     .await
-                    .expect("stake table setup failed");
+                    .expect("stake table setup failed")
+                    .apply_all()
+                    .await
+                    .expect("send all txns failed");
 
                     Upgrade::pos_view_based(st_addr)
                 },
@@ -1051,7 +1054,9 @@ pub mod testing {
                 drb_upgrade_difficulty: 20,
             };
 
-            let anvil = Anvil::new().args(["--slots-in-an-epoch", "0"]).spawn();
+            let anvil = Anvil::new()
+                .args(["--slots-in-an-epoch", "0", "--balance", "1000000"])
+                .spawn();
 
             let l1_client = L1Client::anvil(&anvil).expect("failed to create l1 client");
             let anvil_provider = AnvilProvider::new(l1_client.provider, Arc::new(anvil));

sequencer/src/persistence.rs

@@ -79,7 +79,7 @@ mod tests {
     };
     use indexmap::IndexMap;
     use portpicker::pick_unused_port;
-    use staking_cli::demo::{setup_stake_table_contract_for_test, DelegationConfig};
+    use staking_cli::demo::{DelegationConfig, StakingTransactions};
     use surf_disco::Client;
     use tide_disco::error::ServerError;
     use tokio::{spawn, time::sleep};
@@ -1438,14 +1438,14 @@ mod tests {
         )
         .unwrap();
 
-        let (_, priv_keys): (Vec<_>, Vec<_>) = (0..200)
+        let (_, priv_keys): (Vec<_>, Vec<_>) = (0..20)
             .map(|i| <PubKey as SignatureKey>::generated_from_seed_indexed([1; 32], i as u64))
             .unzip();
-        let state_key_pairs = (0..200)
+        let state_key_pairs = (0..20)
             .map(|i| StateKeyPair::generate_from_seed_indexed([2; 32], i as u64))
             .collect::<Vec<_>>();
 
-        let validators = staking_priv_keys(&priv_keys, &state_key_pairs, 1000);
+        let validators = staking_priv_keys(&priv_keys, &state_key_pairs, 20);
 
         let deployer = ProviderBuilder::new()
             .wallet(EthereumWallet::from(network_config.signer().clone()))
@@ -1486,6 +1486,24 @@ mod tests {
             .expect("StakeTableProxy deployed");
         let l1_url = network_config.l1_url().clone();
 
+        let mut planned_txns = StakingTransactions::create(
+            l1_url.clone(),
+            &deployer,
+            st_addr,
+            validators,
+            DelegationConfig::MultipleDelegators,
+        )
+        .await
+        .expect("stake table setup failed");
+
+        planned_txns
+            .apply_prerequisites()
+            .await
+            .expect("prerequisites failed");
+
+        // Ensure we have at least one stake table affecting transaction
+        planned_txns.apply_one().await.expect("send tx failed");
+
         // new block every 1s
         anvil_provider
             .anvil_set_interval_mining(1)
@@ -1496,18 +1514,13 @@ mod tests {
         // this is going to keep registering validators and multiple delegators
         // the interval mining is set to 1s so each transaction finalization would take atleast 1s
         spawn({
-            let l1_url = l1_url.clone();
             async move {
                 {
-                    setup_stake_table_contract_for_test(
-                        l1_url,
-                        &deployer,
-                        st_addr,
-                        validators,
-                        DelegationConfig::MultipleDelegators,
-                    )
-                    .await
-                    .expect("stake table setup failed");
+                    while let Some(receipt) =
+                        planned_txns.apply_one().await.expect("send tx failed")
+                    {
+                        tracing::debug!(?receipt, "transaction finalized");
+                    }
                 }
             }
         });

sequencer/src/restart_tests.rs

@@ -58,7 +58,7 @@ use itertools::Itertools;
 use options::Modules;
 use portpicker::pick_unused_port;
 use run::init_with_storage;
-use staking_cli::demo::{setup_stake_table_contract_for_test, DelegationConfig};
+use staking_cli::demo::{DelegationConfig, StakingTransactions};
 use surf_disco::{error::ClientError, Url};
 use tempfile::TempDir;
 use tokio::{
@@ -867,15 +867,18 @@ impl TestNetwork {
 
         tracing::info!(?stake_table_address, ?token_addr);
 
-        setup_stake_table_contract_for_test(
+        StakingTransactions::create(
             l1_url.clone(),
             &deployer,
             stake_table_address,
             staking_priv_keys,
             delegation_config,
         )
         .await
-        .expect("stake table setup failed");
+        .expect("stake table setup failed")
+        .apply_all()
+        .await
+        .expect("send all txns failed");
 
         self.anvil
             .anvil_set_interval_mining(1)

staking-cli/Cargo.toml

@@ -44,6 +44,7 @@ url = { workspace = true }
 [dev-dependencies]
 assert_cmd = "2.0.17"
 predicates = "3.1.3"
+pretty_assertions.workspace = true
 tempfile = { workspace = true }
 test-log = { workspace = true }
 

staking-cli/src/claim.rs

@@ -1,4 +1,8 @@
-use alloy::{primitives::Address, providers::Provider, rpc::types::TransactionReceipt};
+use alloy::{
+    network::Ethereum,
+    primitives::Address,
+    providers::{PendingTransactionBuilder, Provider},
+};
 use anyhow::Result;
 use hotshot_contract_adapter::{
     evm::DecodeRevert as _,
@@ -9,39 +13,32 @@ pub async fn claim_withdrawal(
     provider: impl Provider,
     stake_table: Address,
     validator_address: Address,
-) -> Result<TransactionReceipt> {
+) -> Result<PendingTransactionBuilder<Ethereum>> {
     let st = StakeTable::new(stake_table, provider);
-    // See if there are any logs
-    Ok(st
-        .claimWithdrawal(validator_address)
+    st.claimWithdrawal(validator_address)
         .send()
         .await
-        .maybe_decode_revert::<StakeTableErrors>()?
-        .get_receipt()
-        .await?)
+        .maybe_decode_revert::<StakeTableErrors>()
 }
 
 pub async fn claim_validator_exit(
     provider: impl Provider,
     stake_table: Address,
     validator_address: Address,
-) -> Result<TransactionReceipt> {
+) -> Result<PendingTransactionBuilder<Ethereum>> {
     let st = StakeTable::new(stake_table, provider);
-    Ok(st
-        .claimValidatorExit(validator_address)
+    st.claimValidatorExit(validator_address)
         .send()
         .await
-        .maybe_decode_revert::<StakeTableErrors>()?
-        .get_receipt()
-        .await?)
+        .maybe_decode_revert::<StakeTableErrors>()
 }
 
 #[cfg(test)]
 mod test {
     use alloy::primitives::U256;
 
     use super::*;
-    use crate::deploy::TestSystem;
+    use crate::{deploy::TestSystem, receipt::ReceiptExt};
 
     #[tokio::test]
     async fn test_claim_withdrawal() -> Result<()> {
@@ -53,9 +50,10 @@ mod test {
         system.warp_to_unlock_time().await?;
 
         let validator_address = system.deployer_address;
-        let receipt =
-            claim_withdrawal(&system.provider, system.stake_table, validator_address).await?;
-        assert!(receipt.status());
+        let receipt = claim_withdrawal(&system.provider, system.stake_table, validator_address)
+            .await?
+            .assert_success()
+            .await?;
 
         let event = receipt.decoded_log::<StakeTable::Withdrawal>().unwrap();
         assert_eq!(event.amount, amount);
@@ -73,9 +71,10 @@ mod test {
         system.warp_to_unlock_time().await?;
 
         let validator_address = system.deployer_address;
-        let receipt =
-            claim_validator_exit(&system.provider, system.stake_table, validator_address).await?;
-        assert!(receipt.status());
+        let receipt = claim_validator_exit(&system.provider, system.stake_table, validator_address)
+            .await?
+            .assert_success()
+            .await?;
 
         let event = receipt.decoded_log::<StakeTable::Withdrawal>().unwrap();
         assert_eq!(event.amount, amount);