From c4d9aa79a9ca1827625a110e1f4e501c4119e31c Mon Sep 17 00:00:00 2001
From: romsters <petriv.roma@gmail.com>
Date: Thu, 12 Dec 2024 00:15:57 +0200
Subject: [PATCH] chore: improve e2e tests

---
 e2e-tests-rust/src/ext.rs              |  9 +++++++--
 e2e-tests-rust/src/provider/testing.rs | 18 ++++++++++++++++++
 e2e-tests-rust/tests/lib.rs            | 21 +++++++++++++++++++--
 3 files changed, 44 insertions(+), 4 deletions(-)

diff --git a/e2e-tests-rust/src/ext.rs b/e2e-tests-rust/src/ext.rs
index 7114a6d8..229fcc3e 100644
--- a/e2e-tests-rust/src/ext.rs
+++ b/e2e-tests-rust/src/ext.rs
@@ -1,7 +1,7 @@
 //! Module containing extensions of existing alloy abstractions.
 
-use alloy::network::ReceiptResponse;
-use alloy::primitives::{Address, BlockHash};
+use alloy::network::{ReceiptResponse,TxSigner};
+use alloy::primitives::{Address,BlockHash,PrimitiveSignature};
 use alloy::providers::WalletProvider;
 use alloy::signers::local::PrivateKeySigner;
 use alloy_zksync::network::Zksync;
@@ -78,6 +78,11 @@ pub trait ZksyncWalletProviderExt: WalletProvider<Zksync, Wallet = ZksyncWallet>
         self.wallet_mut().register_signer(signer);
         address
     }
+
+    /// Registers provider signer.
+    fn register_signer<T: TxSigner<PrimitiveSignature> + Send + Sync + 'static>(&mut self, signer: T) {
+        self.wallet_mut().register_signer(signer);
+    }
 }
 
 impl<T: WalletProvider<Zksync, Wallet = ZksyncWallet>> ZksyncWalletProviderExt for T {}
diff --git a/e2e-tests-rust/src/provider/testing.rs b/e2e-tests-rust/src/provider/testing.rs
index dc5635d2..48e713b6 100644
--- a/e2e-tests-rust/src/provider/testing.rs
+++ b/e2e-tests-rust/src/provider/testing.rs
@@ -383,6 +383,24 @@ where
         self
     }
 
+    /// Builder-pattern method for setting the chain id.
+    pub fn with_chain_id(mut self, id: u64) -> Self {
+        self.inner = self.inner.with_chain_id(id);
+        self
+    }
+
+    /// Builder-pattern method for setting the recipient.
+    pub fn with_to(mut self, to: Address) -> Self {
+        self.inner = self.inner.with_to(to);
+        self
+    }
+
+    /// Builder-pattern method for setting the value.
+    pub fn with_value(mut self, value: U256) -> Self {
+        self.inner = self.inner.with_value(value);
+        self
+    }
+
     /// Submits transaction to the node.
     ///
     /// This does not wait for the transaction to be confirmed, but returns a [`PendingTransactionFinalizable`]
diff --git a/e2e-tests-rust/tests/lib.rs b/e2e-tests-rust/tests/lib.rs
index 7212db1d..05a4f300 100644
--- a/e2e-tests-rust/tests/lib.rs
+++ b/e2e-tests-rust/tests/lib.rs
@@ -4,6 +4,10 @@ use alloy::providers::Provider;
 use anvil_zksync_e2e_tests::{
     init_testing_provider, AnvilZKsyncApi, ReceiptExt, ZksyncWalletProviderExt, DEFAULT_TX_VALUE,
 };
+use alloy::{
+    primitives::U256,
+    signers::local::PrivateKeySigner,
+};
 use std::convert::identity;
 use std::time::Duration;
 
@@ -337,12 +341,25 @@ async fn cant_load_into_existing_state() -> anyhow::Result<()> {
 
 #[tokio::test]
 async fn set_chain_id() -> anyhow::Result<()> {
-    let provider = init_testing_provider(|node| node.block_time(3)).await?;
+    let mut provider = init_testing_provider(identity).await?;
+
+    let random_signer = PrivateKeySigner::random();
+    let random_signer_address = random_signer.address();
 
-    let new_chain_id = 261;
+    // Send transaction before changing chain id
+    provider.tx().with_to(random_signer_address).with_value(U256::from(1e18 as u64)).finalize().await?;
+
+    // Change chain id
+    let new_chain_id = 123;
     provider.anvil_set_chain_id(new_chain_id).await?;
 
+    // Verify new chain id
     assert_eq!(new_chain_id, provider.get_chain_id().await?);
 
+    // Verify transactions can be executed after chain id change
+    // Registering and using new signer to get new chain id applied
+    provider.register_signer(random_signer);
+    provider.tx().with_from(random_signer_address).with_chain_id(new_chain_id).finalize().await?;
+
     Ok(())
 }