Skip to content

Commit

Permalink
feat(cheatcodes): vm.getScriptWallets() (#9052)
Browse files Browse the repository at this point in the history
* feat(`cheatcodes`): vm.getScriptWallets()

* feat: load default anvil accounts in script

* Revert "feat: load default anvil accounts in script"

This reverts commit 4d64356.

* clippy

* test

---------

Co-authored-by: grandizzy <grandizzy.the.egg@gmail.com>

---------

Co-authored-by: grandizzy <grandizzy.the.egg@gmail.com>
  • Loading branch information
yash-atreya and grandizzy authored Oct 9, 2024
1 parent a17869a commit 373ad46
Show file tree
Hide file tree
Showing 5 changed files with 80 additions and 0 deletions.
20 changes: 20 additions & 0 deletions crates/cheatcodes/assets/cheatcodes.json

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

4 changes: 4 additions & 0 deletions crates/cheatcodes/spec/src/vm.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 --------
Expand Down
15 changes: 15 additions & 0 deletions crates/cheatcodes/src/script.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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
Expand Down
40 changes: 40 additions & 0 deletions crates/forge/tests/cli/script.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
"#]]);
});
1 change: 1 addition & 0 deletions testdata/cheats/Vm.sol

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

0 comments on commit 373ad46

Please sign in to comment.