foundry_test_utils/
script.rs

1use crate::{TestCommand, init_tracing, util::lossy_string};
2use alloy_primitives::{Address, address};
3use alloy_provider::Provider;
4use eyre::Result;
5use foundry_common::provider::{RetryProvider, get_http_provider};
6use std::{
7    collections::BTreeMap,
8    fs,
9    path::{Path, PathBuf},
10};
11
12const BROADCAST_TEST_PATH: &str = "src/Broadcast.t.sol";
13const TESTDATA: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/../../testdata");
14
15fn init_script_cmd(
16    cmd: &mut TestCommand,
17    project_root: &Path,
18    target_contract: &str,
19    endpoint: Option<&str>,
20) {
21    cmd.forge_fuse();
22    cmd.set_current_dir(project_root);
23
24    cmd.args(["script", target_contract, "--root", project_root.to_str().unwrap(), "-vvvvv"]);
25
26    if let Some(rpc_url) = endpoint {
27        cmd.args(["--fork-url", rpc_url]);
28    }
29}
30/// A helper struct to test forge script scenarios
31pub struct ScriptTester {
32    pub accounts_pub: Vec<Address>,
33    pub accounts_priv: Vec<String>,
34    pub provider: Option<RetryProvider>,
35    pub nonces: BTreeMap<u32, u64>,
36    pub address_nonces: BTreeMap<Address, u64>,
37    pub cmd: TestCommand,
38    pub project_root: PathBuf,
39    pub target_contract: String,
40    pub endpoint: Option<String>,
41}
42
43impl ScriptTester {
44    /// Creates a new instance of a Tester for the given contract
45    pub fn new(
46        mut cmd: TestCommand,
47        endpoint: Option<&str>,
48        project_root: &Path,
49        target_contract: &str,
50    ) -> Self {
51        init_tracing();
52        Self::copy_testdata(project_root).unwrap();
53        init_script_cmd(&mut cmd, project_root, target_contract, endpoint);
54
55        let mut provider = None;
56        if let Some(endpoint) = endpoint {
57            provider = Some(get_http_provider(endpoint))
58        }
59
60        Self {
61            accounts_pub: vec![
62                address!("0xf39Fd6e51aad88F6F4ce6aB8827279cffFb92266"),
63                address!("0x70997970C51812dc3A010C7d01b50e0d17dc79C8"),
64                address!("0x3C44CdDdB6a900fa2b585dd299e03d12FA4293BC"),
65            ],
66            accounts_priv: vec![
67                "ac0974bec39a17e36ba4a6b4d238ff944bacb478cbed5efcae784d7bf4f2ff80".to_string(),
68                "59c6995e998f97a5a0044966f0945389dc9e86dae88c7a8412f4603b6b78690d".to_string(),
69                "5de4111afa1a4b94908f83103eb1f1706367c2e68ca870fc3fb9a804cdab365a".to_string(),
70            ],
71            provider,
72            nonces: BTreeMap::default(),
73            address_nonces: BTreeMap::default(),
74            cmd,
75            project_root: project_root.to_path_buf(),
76            target_contract: target_contract.to_string(),
77            endpoint: endpoint.map(|s| s.to_string()),
78        }
79    }
80
81    /// Creates a new instance of a Tester for the `broadcast` test at the given `project_root` by
82    /// configuring the `TestCommand` with script
83    pub fn new_broadcast(cmd: TestCommand, endpoint: &str, project_root: &Path) -> Self {
84        let target_contract = project_root.join(BROADCAST_TEST_PATH).to_string_lossy().to_string();
85
86        // copy the broadcast test
87        fs::copy(
88            Self::testdata_path().join("default/cheats/Broadcast.t.sol"),
89            project_root.join(BROADCAST_TEST_PATH),
90        )
91        .expect("Failed to initialize broadcast contract");
92
93        Self::new(cmd, Some(endpoint), project_root, &target_contract)
94    }
95
96    /// Creates a new instance of a Tester for the `broadcast` test at the given `project_root` by
97    /// configuring the `TestCommand` with script without an endpoint
98    pub fn new_broadcast_without_endpoint(cmd: TestCommand, project_root: &Path) -> Self {
99        let target_contract = project_root.join(BROADCAST_TEST_PATH).to_string_lossy().to_string();
100
101        // copy the broadcast test
102        let testdata = Self::testdata_path();
103        fs::copy(
104            testdata.join("default/cheats/Broadcast.t.sol"),
105            project_root.join(BROADCAST_TEST_PATH),
106        )
107        .expect("Failed to initialize broadcast contract");
108
109        Self::new(cmd, None, project_root, &target_contract)
110    }
111
112    /// Returns the path to the dir that contains testdata
113    fn testdata_path() -> &'static Path {
114        Path::new(TESTDATA)
115    }
116
117    /// Initialises the test contracts by copying them into the workspace
118    fn copy_testdata(root: &Path) -> Result<()> {
119        let testdata = Self::testdata_path();
120        let from_dir = testdata.join("utils");
121        let to_dir = root.join("utils");
122        fs::create_dir_all(&to_dir)?;
123        for entry in fs::read_dir(&from_dir)? {
124            let file = &entry?.path();
125            let name = file.file_name().unwrap();
126            fs::copy(file, to_dir.join(name))?;
127        }
128        Ok(())
129    }
130
131    pub async fn load_private_keys(&mut self, keys_indexes: &[u32]) -> &mut Self {
132        for &index in keys_indexes {
133            self.cmd.args(["--private-keys", &self.accounts_priv[index as usize]]);
134
135            if let Some(provider) = &self.provider {
136                let nonce = provider
137                    .get_transaction_count(self.accounts_pub[index as usize])
138                    .await
139                    .unwrap();
140                self.nonces.insert(index, nonce);
141            }
142        }
143        self
144    }
145
146    pub async fn load_addresses(&mut self, addresses: &[Address]) -> &mut Self {
147        for &address in addresses {
148            let nonce =
149                self.provider.as_ref().unwrap().get_transaction_count(address).await.unwrap();
150            self.address_nonces.insert(address, nonce);
151        }
152        self
153    }
154
155    pub fn add_deployer(&mut self, index: u32) -> &mut Self {
156        self.sender(self.accounts_pub[index as usize])
157    }
158
159    /// Adds given address as sender
160    pub fn sender(&mut self, addr: Address) -> &mut Self {
161        self.args(&["--sender", &addr.to_string()])
162    }
163
164    pub fn add_sig(&mut self, contract_name: &str, sig: &str) -> &mut Self {
165        self.args(&["--tc", contract_name, "--sig", sig])
166    }
167
168    pub fn add_create2_deployer(&mut self, create2_deployer: Address) -> &mut Self {
169        self.args(&["--create2-deployer", &create2_deployer.to_string()])
170    }
171
172    /// Adds the `--unlocked` flag
173    pub fn unlocked(&mut self) -> &mut Self {
174        self.arg("--unlocked")
175    }
176
177    pub fn simulate(&mut self, expected: ScriptOutcome) -> &mut Self {
178        self.run(expected)
179    }
180
181    pub fn broadcast(&mut self, expected: ScriptOutcome) -> &mut Self {
182        self.arg("--broadcast").run(expected)
183    }
184
185    pub fn resume(&mut self, expected: ScriptOutcome) -> &mut Self {
186        self.arg("--resume").run(expected)
187    }
188
189    /// `[(private_key_slot, expected increment)]`
190    pub async fn assert_nonce_increment(&mut self, keys_indexes: &[(u32, u32)]) -> &mut Self {
191        for &(private_key_slot, expected_increment) in keys_indexes {
192            let addr = self.accounts_pub[private_key_slot as usize];
193            let nonce = self.provider.as_ref().unwrap().get_transaction_count(addr).await.unwrap();
194            let prev_nonce = self.nonces.get(&private_key_slot).unwrap();
195
196            assert_eq!(
197                nonce,
198                (*prev_nonce + expected_increment as u64),
199                "nonce not incremented correctly for {addr}: \
200                 {prev_nonce} + {expected_increment} != {nonce}"
201            );
202        }
203        self
204    }
205
206    /// In Vec<(address, expected increment)>
207    pub async fn assert_nonce_increment_addresses(
208        &mut self,
209        address_indexes: &[(Address, u32)],
210    ) -> &mut Self {
211        for (address, expected_increment) in address_indexes {
212            let nonce =
213                self.provider.as_ref().unwrap().get_transaction_count(*address).await.unwrap();
214            let prev_nonce = self.address_nonces.get(address).unwrap();
215
216            assert_eq!(nonce, *prev_nonce + *expected_increment as u64);
217        }
218        self
219    }
220
221    pub fn run(&mut self, expected: ScriptOutcome) -> &mut Self {
222        let out = self.cmd.execute();
223        let (stdout, stderr) = (lossy_string(&out.stdout), lossy_string(&out.stderr));
224
225        trace!(target: "tests", "STDOUT\n{stdout}\n\nSTDERR\n{stderr}");
226
227        if !stdout.contains(expected.as_str()) && !stderr.contains(expected.as_str()) {
228            panic!(
229                "--STDOUT--\n{stdout}\n\n--STDERR--\n{stderr}\n\n--EXPECTED--\n{:?} not found in stdout or stderr",
230                expected.as_str()
231            );
232        }
233
234        self
235    }
236
237    pub fn slow(&mut self) -> &mut Self {
238        self.arg("--slow")
239    }
240
241    pub fn arg(&mut self, arg: &str) -> &mut Self {
242        self.cmd.arg(arg);
243        self
244    }
245
246    pub fn args(&mut self, args: &[&str]) -> &mut Self {
247        self.cmd.args(args);
248        self
249    }
250
251    pub fn clear(&mut self) {
252        init_script_cmd(
253            &mut self.cmd,
254            &self.project_root,
255            &self.target_contract,
256            self.endpoint.as_deref(),
257        );
258        self.nonces.clear();
259        self.address_nonces.clear();
260    }
261}
262
263/// Various `forge` script results
264#[derive(Debug)]
265pub enum ScriptOutcome {
266    OkNoEndpoint,
267    OkSimulation,
268    OkBroadcast,
269    WarnSpecifyDeployer,
270    MissingSender,
271    MissingWallet,
272    StaticCallNotAllowed,
273    ScriptFailed,
274    UnsupportedLibraries,
275    ErrorSelectForkOnBroadcast,
276    OkRun,
277}
278
279impl ScriptOutcome {
280    pub fn as_str(&self) -> &'static str {
281        match self {
282            Self::OkNoEndpoint => "If you wish to simulate on-chain transactions pass a RPC URL.",
283            Self::OkSimulation => "SIMULATION COMPLETE. To broadcast these",
284            Self::OkBroadcast => "ONCHAIN EXECUTION COMPLETE & SUCCESSFUL",
285            Self::WarnSpecifyDeployer => {
286                "Warning: You have more than one deployer who could predeploy libraries. Using `--sender` instead."
287            }
288            Self::MissingSender => {
289                "You seem to be using Foundry's default sender. Be sure to set your own --sender"
290            }
291            Self::MissingWallet => "No associated wallet",
292            Self::StaticCallNotAllowed => {
293                "staticcall`s are not allowed after `broadcast`; use `startBroadcast` instead"
294            }
295            Self::ScriptFailed => "script failed: ",
296            Self::UnsupportedLibraries => {
297                "Multi chain deployment does not support library linking at the moment."
298            }
299            Self::ErrorSelectForkOnBroadcast => "cannot select forks during a broadcast",
300            Self::OkRun => "Script ran successfully",
301        }
302    }
303
304    pub fn is_err(&self) -> bool {
305        match self {
306            Self::OkNoEndpoint
307            | Self::OkSimulation
308            | Self::OkBroadcast
309            | Self::WarnSpecifyDeployer
310            | Self::OkRun => false,
311            Self::MissingSender
312            | Self::MissingWallet
313            | Self::StaticCallNotAllowed
314            | Self::UnsupportedLibraries
315            | Self::ErrorSelectForkOnBroadcast
316            | Self::ScriptFailed => true,
317        }
318    }
319}
foundry_test_utils/script.rs

foundry_test_utils/
script.rs
