huff_neo_test_runner/
lib.rs

1use crate::{errors::RunnerError, types::TestResult};
2use huff_neo_utils::prelude::{Contract, MacroDefinition, SupportedEVMVersions};
3
4/// The runner module
5pub mod runner;
6
7/// The report module
8pub mod report;
9
10/// The types module
11pub mod types;
12
13/// The errors module
14pub mod errors;
15
16/// The decorator spec module
17pub mod spec;
18
19/// The test outcome module
20pub mod outcome;
21
22/// The module that shows Huff functions as internal calls in traces
23pub mod internal;
24
25use alloy_primitives::map::AddressHashMap;
26use foundry_compilers::{compilers::solc::SolcLanguage, multi::MultiCompilerLanguage};
27use foundry_debugger::Debugger;
28use foundry_evm::traces::TraceKind;
29use foundry_evm_traces::debug::{ArtifactData, ContractSources, SourceData};
30use revm::primitives::hardfork::SpecId;
31use std::path::PathBuf;
32use std::sync::Arc;
33
34pub use crate::runner::{Env, TestRunner, cheats_config};
35pub use foundry_evm::inspectors::CheatsConfig;
36pub use revm::context::TxEnv;
37
38/// Prelude wraps all modules within the crate
39pub mod prelude {
40    pub use crate::{errors::*, internal::*, outcome::*, report::*, runner::*, spec::*, types::*};
41}
42
43/// A vector of shared references to test macro definitions
44pub type TestMacros<'t> = Vec<&'t MacroDefinition>;
45
46/// Returns the REVM hardfork that matches a Huff EVM version
47pub fn spec_id_for(version: &SupportedEVMVersions) -> SpecId {
48    match version {
49        SupportedEVMVersions::Paris => SpecId::MERGE,
50        SupportedEVMVersions::Shanghai => SpecId::SHANGHAI,
51        SupportedEVMVersions::Cancun => SpecId::CANCUN,
52        SupportedEVMVersions::Prague => SpecId::PRAGUE,
53        SupportedEVMVersions::Osaka => SpecId::OSAKA,
54        SupportedEVMVersions::Amsterdam => SpecId::AMSTERDAM,
55    }
56}
57
58/// The core struct of the huff-tests crate.
59///
60/// A `HuffTester` struct is instantiated with an AST of a contract that contains test
61/// macros. The struct can be consumed by the [HuffTester::execute] method,
62/// returning a vector of [TestResult] structs.
63pub struct HuffTester<'t> {
64    /// The AST of the contract
65    pub ast: &'t Contract,
66
67    /// The test macros
68    pub macros: TestMacros<'t>,
69
70    /// The test runner
71    pub runner: TestRunner,
72
73    /// Whether to open the single matched test in the debugger after running it
74    pub debug: bool,
75}
76
77/// HuffTester implementation
78impl<'t> HuffTester<'t> {
79    /// Create a new instance of `HuffTester` from a contract's AST.
80    pub fn new(ast: &'t Contract, match_test_name: Option<String>, runner: TestRunner, debug: bool) -> Self {
81        Self {
82            ast,
83            macros: {
84                // Filter all macros within the AST for `test` macros only
85                let mut macros: TestMacros<'t> = ast.macros.values().filter(|m| m.test).collect();
86                // If the match flag is present, only retain the test macro
87                // that was queried
88                if let Some(match_test_name) = match_test_name {
89                    macros.retain(|m| m.name == *match_test_name);
90                }
91                macros
92            },
93            runner,
94            debug,
95        }
96    }
97
98    /// Execute tests
99    pub fn execute(self) -> Result<Vec<TestResult>, RunnerError> {
100        // Check if any test macros exist
101        if self.macros.is_empty() {
102            return Err(RunnerError::GenericError(String::from("No test macros found.")));
103        }
104
105        // Execute our tests and return a vector of the results
106        let results = self
107            .macros
108            .into_iter()
109            .map(|macro_def| {
110                // Every test gets its own executor; each one must fork the same chain state
111                let mut executor = self.runner.new_executor()?;
112                self.runner.run_test(&mut executor, macro_def, self.ast)
113            })
114            .collect::<Result<Vec<TestResult>, RunnerError>>();
115
116        if !self.debug {
117            return results;
118        }
119
120        //
121        // Run the debugger
122        //
123        let Ok(ref results_vec) = results else {
124            return results;
125        };
126        if results_vec.len() != 1 {
127            return Err(RunnerError::GenericError("Debugging is only supported for a single test".to_string()));
128        }
129        let Some(test_result) = results_vec.first() else {
130            return Err(RunnerError::GenericError("No test result found".to_string()));
131        };
132
133        let Some(arena) = test_result.traces.clone() else {
134            return Err(RunnerError::GenericError("No trace found".to_string()));
135        };
136
137        let traces = [(TraceKind::Execution, arena)];
138
139        eprintln!("DEBUG: Starting debugger for test: {}", test_result.name);
140        eprintln!("DEBUG: Has source map: {}", test_result.source_map.is_some());
141        eprintln!("DEBUG: Has source code: {}", test_result.source_code.is_some());
142        if let Some(addr) = test_result.address {
143            eprintln!("DEBUG: Contract address: {}", addr);
144        }
145
146        // Create ContractSources if we have source maps
147        let mut debugger_builder = Debugger::builder().traces(traces.iter().filter(|(t, _)| t.is_execution()).cloned().collect());
148
149        // Add identified contracts and source maps for debugging
150        let mut identified_contracts = AddressHashMap::default();
151        let mut contracts_sources = ContractSources::default();
152
153        if let Some(address) = test_result.address {
154            let contract_name = format!("Test::{}", test_result.name);
155            identified_contracts.insert(address, contract_name.clone());
156
157            // Add source maps if available
158            if let Some(source_map) = &test_result.source_map {
159                // Create inner map for file IDs
160                let mut file_map = std::collections::HashMap::new();
161
162                // Check if we have multiple source files
163                if !test_result.source_files.is_empty() {
164                    // Add each source file with its proper file ID
165                    for (file_id, (file_path, source_content)) in test_result.source_files.iter().enumerate() {
166                        let file_name = PathBuf::from(file_path)
167                            .file_stem()
168                            .map(|s| s.to_string_lossy().to_string())
169                            .unwrap_or_else(|| format!("file_{}", file_id));
170
171                        let source_data = SourceData {
172                            source: Arc::new(source_content.clone()),
173                            language: MultiCompilerLanguage::Solc(SolcLanguage::Solidity),
174                            path: PathBuf::from(file_path),
175                            contract_definitions: vec![(file_name, 0..source_content.len())],
176                            debug_scopes: Vec::new(),
177                        };
178                        file_map.insert(file_id as u32, Arc::new(source_data));
179                    }
180                } else {
181                    // Fallback to single flattened source
182                    let source_code = test_result
183                        .source_code
184                        .clone()
185                        .unwrap_or_else(|| format!("// Huff test macro: {}\n// Source not available", test_result.name));
186
187                    let source_data = SourceData {
188                        source: Arc::new(source_code.clone()),
189                        language: MultiCompilerLanguage::Solc(SolcLanguage::Solidity),
190                        path: PathBuf::from("test.huff"),
191                        contract_definitions: vec![("test".to_string(), 0..source_code.len())],
192                        debug_scopes: Vec::new(),
193                    };
194                    file_map.insert(0u32, Arc::new(source_data));
195                }
196
197                // Insert into sources_by_id with address as string key
198                contracts_sources.sources_by_id.insert(address.to_string(), file_map);
199
200                // Also need to add the artifact data with the source map
201                let solidity_source_map = huff_neo_utils::artifact::SourceMapEntry::generate_solidity_source_map(source_map);
202
203                // Parse the source map
204                let parsed_source_map = foundry_compilers::artifacts::sourcemap::parse(&solidity_source_map).ok();
205
206                if let Some(source_map_parsed) = parsed_source_map {
207                    // Create PcIcMap from bytecode if available
208                    let pc_ic_map = test_result.bytecode.as_ref().and_then(|bytecode_hex| {
209                        alloy_primitives::hex::decode(bytecode_hex).ok().map(|bytes| foundry_evm_core::ic::PcIcMap::new(&bytes))
210                    });
211
212                    // Add artifact data for the contract
213                    // The build_id should match the key in sources_by_id
214                    let artifact_data = ArtifactData {
215                        build_id: address.to_string(),
216                        file_id: 0,
217                        source_map: Some(source_map_parsed.clone()),
218                        source_map_runtime: Some(source_map_parsed),
219                        pc_ic_map: pc_ic_map.clone(),
220                        pc_ic_map_runtime: pc_ic_map,
221                    };
222
223                    // Try adding artifacts by both name AND address
224                    // Add by contract name
225                    let contract_name = format!("Test::{}", test_result.name);
226                    contracts_sources.artifacts_by_name.entry(contract_name).or_default().push(artifact_data.clone());
227
228                    // Also add by address in case debugger looks it up that way
229                    contracts_sources.artifacts_by_name.entry(address.to_string()).or_default().push(artifact_data);
230                }
231            }
232        }
233
234        debugger_builder = debugger_builder.identified_contracts(identified_contracts).sources(contracts_sources);
235
236        let mut debugger = debugger_builder.build();
237
238        debugger.try_run_tui().map_err(|e| RunnerError::GenericError(format!("{e:?}")))?;
239
240        results
241    }
242}
huff_neo_test_runner/lib.rs

huff_neo_test_runner/
lib.rs
