From 373ad46de9034f3b9e30b95084c9d1bd076d66a7 Mon Sep 17 00:00:00 2001
From: Yash Atreya <44857776+yash-atreya@users.noreply.github.com>
Date: Wed, 9 Oct 2024 11:31:38 +0530
Subject: [PATCH] feat(`cheatcodes`): vm.getScriptWallets() (#9052)
* feat(`cheatcodes`): vm.getScriptWallets()
* feat: load default anvil accounts in script
* Revert "feat: load default anvil accounts in script"
This reverts commit 4d64356a51bf226482269a2af47f947c4e49e462.
* clippy
* test
---------
Co-authored-by: grandizzy <grandizzy.the.egg@gmail.com>
---------
Co-authored-by: grandizzy <grandizzy.the.egg@gmail.com>
---
crates/cheatcodes/assets/cheatcodes.json | 20 ++++++++++++
crates/cheatcodes/spec/src/vm.rs | 4 +++
crates/cheatcodes/src/script.rs | 15 +++++++++
crates/forge/tests/cli/script.rs | 40 ++++++++++++++++++++++++
testdata/cheats/Vm.sol | 1 +
5 files changed, 80 insertions(+)
diff --git a/crates/cheatcodes/assets/cheatcodes.json b/crates/cheatcodes/assets/cheatcodes.json
index e6fef48b0..06ddba5bd 100644
--- a/crates/cheatcodes/assets/cheatcodes.json
+++ b/crates/cheatcodes/assets/cheatcodes.json
@@ -5395,6 +5395,26 @@
"status": "stable",
"safety": "safe"
},
+ {
+ "func": {
+ "id": "getScriptWallets",
+ "description": "Returns addresses of available unlocked wallets in the script environment.",
+ "declaration": "function getScriptWallets() external returns (address[] memory wallets);",
+ "visibility": "external",
+ "mutability": "",
+ "signature": "getScriptWallets()",
+ "selector": "0x7c49aa1f",
+ "selectorBytes": [
+ 124,
+ 73,
+ 170,
+ 31
+ ]
+ },
+ "group": "scripting",
+ "status": "stable",
+ "safety": "safe"
+ },
{
"func": {
"id": "indexOf",
diff --git a/crates/cheatcodes/spec/src/vm.rs b/crates/cheatcodes/spec/src/vm.rs
index 0ee95e43f..e73755de1 100644
--- a/crates/cheatcodes/spec/src/vm.rs
+++ b/crates/cheatcodes/spec/src/vm.rs
@@ -1877,6 +1877,10 @@ interface Vm {
#[cheatcode(group = Scripting)]
function broadcastRawTransaction(bytes calldata data) external;
+ /// Returns addresses of available unlocked wallets in the script environment.
+ #[cheatcode(group = Scripting)]
+ function getScriptWallets() external returns (address[] memory wallets);
+
// ======== Utilities ========
// -------- Strings --------
diff --git a/crates/cheatcodes/src/script.rs b/crates/cheatcodes/src/script.rs
index 93d5aaaf8..f9535844c 100644
--- a/crates/cheatcodes/src/script.rs
+++ b/crates/cheatcodes/src/script.rs
@@ -3,6 +3,7 @@
use crate::{Cheatcode, CheatsCtxt, Result, Vm::*};
use alloy_primitives::{Address, B256, U256};
use alloy_signer_local::PrivateKeySigner;
+use alloy_sol_types::SolValue;
use foundry_wallets::{multi_wallet::MultiWallet, WalletSigner};
use parking_lot::Mutex;
use std::sync::Arc;
@@ -60,6 +61,20 @@ impl Cheatcode for stopBroadcastCall {
}
}
+impl Cheatcode for getScriptWalletsCall {
+ fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
+ let script_wallets =
+ ccx.state.script_wallets().cloned().map(|sw| sw.signers().unwrap_or_default());
+
+ if let Some(script_wallets) = script_wallets {
+ let script_wallets: Vec<Address> = script_wallets.into_iter().collect();
+ Ok(script_wallets.abi_encode())
+ } else {
+ Ok(Default::default())
+ }
+ }
+}
+
#[derive(Clone, Debug, Default)]
pub struct Broadcast {
/// Address of the transaction origin
diff --git a/crates/forge/tests/cli/script.rs b/crates/forge/tests/cli/script.rs
index 61c2456fe..826d66282 100644
--- a/crates/forge/tests/cli/script.rs
+++ b/crates/forge/tests/cli/script.rs
@@ -2055,3 +2055,43 @@ ONCHAIN EXECUTION COMPLETE & SUCCESSFUL.
"#]]);
});
+
+forgetest_init!(can_get_script_wallets, |prj, cmd| {
+ let script = prj
+ .add_source(
+ "Foo",
+ r#"
+import "forge-std/Script.sol";
+
+interface Vm {
+ function getScriptWallets() external returns (address[] memory wallets);
+}
+
+contract WalletScript is Script {
+ function run() public {
+ address[] memory wallets = Vm(address(vm)).getScriptWallets();
+ console.log(wallets[0]);
+ }
+}"#,
+ )
+ .unwrap();
+ cmd.arg("script")
+ .arg(script)
+ .args([
+ "--private-key",
+ "0x2a871d0798f97d79848a013d4936a73bf4cc922c825d33c1cf7073dff6d409c6",
+ "-v",
+ ])
+ .assert_success()
+ .stdout_eq(str![[r#"
+[COMPILING_FILES] with [SOLC_VERSION]
+[SOLC_VERSION] [ELAPSED]
+Compiler run successful!
+Script ran successfully.
+[GAS]
+
+== Logs ==
+ 0xa0Ee7A142d267C1f36714E4a8F75612F20a79720
+
+"#]]);
+});
diff --git a/testdata/cheats/Vm.sol b/testdata/cheats/Vm.sol
index be6522ac4..1458e3e46 100644
--- a/testdata/cheats/Vm.sol
+++ b/testdata/cheats/Vm.sol
@@ -265,6 +265,7 @@ interface Vm {
function getNonce(address account) external view returns (uint64 nonce);
function getNonce(Wallet calldata wallet) external returns (uint64 nonce);
function getRecordedLogs() external returns (Log[] memory logs);
+ function getScriptWallets() external returns (address[] memory wallets);
function indexOf(string calldata input, string calldata key) external pure returns (uint256);
function isContext(ForgeContext context) external view returns (bool result);
function isDir(string calldata path) external returns (bool result);