Add support overriding block id/tag/hash with a named fork binding

CHANGELOG.md

@@ -22,6 +22,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 #### Added
 
 - Project generated by `snforge` contains `assert_macros` dependency with version 0.1.0 for Scarb <= 2.8.0, otherwise equal to Scarb
+- Support for overriding fork configuration in test attribute with a different block ID, tag, or hash.
 
 ## [0.30.0] - 2024-09-04
 

crates/cheatnet/src/runtime_extensions/forge_config_extension/config.rs

@@ -25,10 +25,17 @@ pub struct InlineForkConfig {
     pub block: BlockId,
 }
 
+#[derive(Debug, Clone, CairoDeserialize, PartialEq)]
+pub struct OverriddenForkConfig {
+    pub name: ByteArray,
+    pub block: BlockId,
+}
+
 #[derive(Debug, Clone, CairoDeserialize, PartialEq)]
 pub enum RawForkConfig {
     Inline(InlineForkConfig),
     Named(ByteArray),
+    Overridden(OverriddenForkConfig),
 }
 
 // fuzzer

crates/forge/src/run_tests/resolve_config.rs

@@ -1,8 +1,9 @@
 use crate::{block_number_map::BlockNumberMap, scarb::config::ForkTarget};
 use anyhow::{anyhow, Result};
 use cheatnet::runtime_extensions::forge_config_extension::config::{
-    BlockId, InlineForkConfig, RawForkConfig,
+    BlockId, InlineForkConfig, OverriddenForkConfig, RawForkConfig,
 };
+use conversions::byte_array::ByteArray;
 use forge_runner::package_tests::{
     with_config::TestTargetWithConfig,
     with_config_resolved::{
@@ -98,21 +99,28 @@ fn parse_block_id(fork_target: &ForkTarget) -> Result<BlockId> {
     Ok(block_id)
 }
 
+fn get_fork_target_from_runner_config<'a>(
+    fork_targets: &'a [ForkTarget],
+    name: &ByteArray,
+) -> Result<&'a ForkTarget> {
+    fork_targets
+        .iter()
+        .find(|fork| fork.name == String::from(name.clone()))
+        .ok_or_else(|| {
+            let name = String::from(name.clone());
+            anyhow!("Fork configuration named = {name} not found in the Scarb.toml")
+        })
+}
+
 fn replace_id_with_params(
     raw_fork_config: RawForkConfig,
     fork_targets: &[ForkTarget],
 ) -> Result<InlineForkConfig> {
     match raw_fork_config {
         RawForkConfig::Inline(raw_fork_params) => Ok(raw_fork_params),
         RawForkConfig::Named(name) => {
-            let fork_target_from_runner_config = fork_targets
-                .iter()
-                .find(|fork| fork.name == String::from(name.clone()))
-                .ok_or_else(|| {
-                    let name = String::from(name);
-
-                    anyhow!("Fork configuration named = {name} not found in the Scarb.toml")
-                })?;
+            let fork_target_from_runner_config =
+                get_fork_target_from_runner_config(fork_targets, &name)?;
 
             let block_id = parse_block_id(fork_target_from_runner_config)?;
 
@@ -121,6 +129,14 @@ fn replace_id_with_params(
                 block: block_id,
             })
         }
+        RawForkConfig::Overridden(OverriddenForkConfig { name, block }) => {
+            let fork_target_from_runner_config =
+                get_fork_target_from_runner_config(fork_targets, &name)?;
+
+            let url = fork_target_from_runner_config.url.parse()?;
+
+            Ok(InlineForkConfig { url, block })
+        }
     }
 }
 

crates/forge/tests/integration/setup_fork.rs

@@ -178,6 +178,97 @@ fn fork_aliased_decorator() {
     assert_passed(&result);
 }
 
+#[test]
+fn fork_aliased_decorator_overrding() {
+    let test = test_case!(indoc!(
+        r#"
+            use starknet::syscalls::get_execution_info_syscall;
+
+            #[test]
+            #[fork("FORK_NAME_FROM_SCARB_TOML", block_number: 2137)]
+            fn test_increase_balance() {
+                let execution_info = get_execution_info_syscall().unwrap().deref();
+                let block_info = execution_info.block_info.deref();
+                let block_number = block_info.block_number;
+
+                assert(block_number == 2137, 'Invalid block');
+            }
+        "#
+    ));
+
+    let rt = Runtime::new().expect("Could not instantiate Runtime");
+
+    ScarbCommand::new_with_stdio()
+        .current_dir(test.path().unwrap())
+        .arg("build")
+        .arg("--test")
+        .run()
+        .unwrap();
+
+    let metadata = ScarbCommand::metadata()
+        .current_dir(test.path().unwrap())
+        .run()
+        .unwrap();
+
+    let package = metadata
+        .packages
+        .iter()
+        .find(|p| p.name == "test_package")
+        .unwrap();
+
+    let raw_test_targets =
+        load_test_artifacts(&test.path().unwrap().join("target/dev"), package).unwrap();
+
+    let result = rt
+        .block_on(run_for_package(
+            RunForPackageArgs {
+                test_targets: raw_test_targets,
+                package_name: "test_package".to_string(),
+                tests_filter: TestsFilter::from_flags(
+                    None,
+                    false,
+                    false,
+                    false,
+                    false,
+                    Default::default(),
+                ),
+                forge_config: Arc::new(ForgeConfig {
+                    test_runner_config: Arc::new(TestRunnerConfig {
+                        exit_first: false,
+                        fuzzer_runs: NonZeroU32::new(256).unwrap(),
+                        fuzzer_seed: 12345,
+                        max_n_steps: None,
+                        is_vm_trace_needed: false,
+                        cache_dir: Utf8PathBuf::from_path_buf(tempdir().unwrap().into_path())
+                            .unwrap()
+                            .join(CACHE_DIR),
+                        contracts_data: ContractsData::try_from(test.contracts().unwrap()).unwrap(),
+                        environment_variables: test.env().clone(),
+                    }),
+                    output_config: Arc::new(OutputConfig {
+                        detailed_resources: false,
+                        execution_data_to_save: ExecutionDataToSave::default(),
+                        versioned_programs_dir: Utf8PathBuf::from_path_buf(
+                            tempdir().unwrap().into_path(),
+                        )
+                        .unwrap()
+                        .join(VERSIONED_PROGRAMS_DIR),
+                    }),
+                }),
+                fork_targets: vec![ForkTarget::new(
+                    "FORK_NAME_FROM_SCARB_TOML".to_string(),
+                    node_rpc_url().to_string(),
+                    "tag".to_string(),
+                    "latest".to_string(),
+                )],
+            },
+            &mut BlockNumberMap::default(),
+        ))
+        .expect("Runner fail");
+
+    assert_passed(&result);
+}
+
 #[test]
 fn fork_cairo0_contract() {
     let test = test_case!(formatdoc!(

crates/snforge-scarb-plugin/src/args.rs

@@ -84,6 +84,11 @@ impl Arguments {
         }
     }
 
+    #[inline]
+    pub fn unnamed(&self) -> UnnamedArgs {
+        UnnamedArgs::new(&self.unnamed)
+    }
+
     #[inline]
     pub fn assert_is_empty<T: AttributeInfo>(&self) -> Result<(), Diagnostic> {
         if self.is_empty() {

crates/snforge-scarb-plugin/src/args/unnamed.rs

@@ -3,10 +3,10 @@ use cairo_lang_macro::Diagnostic;
 use cairo_lang_syntax::node::ast::Expr;
 use std::{collections::HashMap, ops::Deref};
 
-pub struct UnnamedArgs<'a>(Vec<&'a Expr>);
+pub struct UnnamedArgs<'a>(Vec<(usize, &'a Expr)>);
 
 impl<'a> Deref for UnnamedArgs<'a> {
-    type Target = Vec<&'a Expr>;
+    type Target = Vec<(usize, &'a Expr)>;
 
     fn deref(&self) -> &Self::Target {
         &self.0
@@ -19,7 +19,7 @@ impl UnnamedArgs<'_> {
 
         args.sort_by(|(a, _), (b, _)| a.cmp(b));
 
-        let args = args.into_iter().map(|(_, expr)| expr).collect();
+        let args = args.into_iter().map(|(&pos, expr)| (pos, expr)).collect();
 
         UnnamedArgs(args)
     }
@@ -28,7 +28,7 @@ impl UnnamedArgs<'_> {
 impl<'a> UnnamedArgs<'a> {
     pub fn of_length<const N: usize, T: AttributeInfo>(
         &self,
-    ) -> Result<&[&'a Expr; N], Diagnostic> {
+    ) -> Result<&[(usize, &'a Expr); N], Diagnostic> {
         self.as_slice()
             .try_into()
             .map_err(|_| T::error(format!("expected {} arguments, got: {}", N, self.len())))

crates/snforge-scarb-plugin/src/attributes/available_gas.rs

@@ -26,7 +26,7 @@ impl AttributeCollector for AvailableGasCollector {
     ) -> Result<String, Diagnostics> {
         let &[arg] = args.unnamed_only::<Self>()?.of_length::<1, Self>()?;
 
-        let gas = Number::parse_from_expr::<Self>(db, arg, "0")?;
+        let gas = Number::parse_from_expr::<Self>(db, arg.1, arg.0.to_string().as_str())?;
 
         let gas = gas.as_cairo_expression();
 

crates/snforge-scarb-plugin/src/attributes/fork.rs

@@ -2,12 +2,12 @@ use self::block_id::{BlockId, BlockIdVariants};
 use crate::{
     args::Arguments,
     attributes::{AttributeCollector, AttributeInfo, AttributeTypeData},
+    branch,
     cairo_expression::CairoExpression,
     config_statement::extend_with_config_cheatcodes,
     types::ParseFromExpr,
-    utils::branch,
 };
-use cairo_lang_macro::{Diagnostic, Diagnostics, ProcMacroResult, TokenStream};
+use cairo_lang_macro::{Diagnostic, Diagnostics, ProcMacroResult, Severity, TokenStream};
 use cairo_lang_syntax::node::db::SyntaxGroup;
 use indoc::formatdoc;
 use url::Url;
@@ -30,7 +30,11 @@ impl AttributeCollector for ForkCollector {
         args: Arguments,
         _warns: &mut Vec<Diagnostic>,
     ) -> Result<String, Diagnostics> {
-        let expr = branch(inline_args(db, &args), || from_file_args(db, &args))?;
+        let expr = branch!(
+            inline_args(db, &args),
+            overridden_args(db, &args),
+            from_file_args(db, &args)
+        )?;
 
         Ok(expr)
     }
@@ -69,7 +73,7 @@ fn from_file_args(db: &dyn SyntaxGroup, args: &Arguments) -> Result<String, Diag
         .unnamed_only::<ForkCollector>()?
         .of_length::<1, ForkCollector>()?;
 
-    let name = String::parse_from_expr::<ForkCollector>(db, arg, "0")?;
+    let name = String::parse_from_expr::<ForkCollector>(db, arg.1, arg.0.to_string().as_str())?;
 
     let name = name.as_cairo_expression();
 
@@ -78,6 +82,31 @@ fn from_file_args(db: &dyn SyntaxGroup, args: &Arguments) -> Result<String, Diag
     ))
 }
 
+fn overridden_args(db: &dyn SyntaxGroup, args: &Arguments) -> Result<String, Diagnostic> {
+    let &[arg] = args.unnamed().of_length::<1, ForkCollector>()?;
+
+    let block_id = args.named.one_of_once(&[
+        BlockIdVariants::Hash,
+        BlockIdVariants::Number,
+        BlockIdVariants::Tag,
+    ])?;
+
+    let block_id = BlockId::parse_from_expr::<ForkCollector>(db, &block_id, block_id.0.as_ref())?;
+    let name = String::parse_from_expr::<ForkCollector>(db, arg.1, arg.0.to_string().as_str())?;
+
+    let block_id = block_id.as_cairo_expression();
+    let name = name.as_cairo_expression();
+
+    Ok(formatdoc!(
+        "
+            snforge_std::_config_types::ForkConfig::Overridden(
+                block: {block_id},
+                name: {name}
+            )
+        "
+    ))
+}
+
 #[must_use]
 pub fn fork(args: TokenStream, item: TokenStream) -> ProcMacroResult {
     extend_with_config_cheatcodes::<ForkCollector>(args, item)

crates/snforge-scarb-plugin/src/utils.rs

@@ -1,30 +1,55 @@
 use cairo_lang_macro::{Diagnostic, Severity};
 use indoc::formatdoc;
 
-fn higher_severity(a: Severity, b: Severity) -> Severity {
+pub fn higher_severity(a: Severity, b: Severity) -> Severity {
     match (a, b) {
         (Severity::Warning, Severity::Warning) => Severity::Warning,
         _ => Severity::Error,
     }
 }
+pub fn format_error_message(variants: &[Diagnostic]) -> String {
+    let formatted_variants: Vec<String> = variants
+        .iter()
+        .map(|variant| format!("- variant: {}", variant.message))
+        .collect();
 
-pub fn branch(
-    left: Result<String, Diagnostic>,
-    right: impl Fn() -> Result<String, Diagnostic>,
-) -> Result<String, Diagnostic> {
-    left.or_else(|error| {
-        right().map_err(|next_error| Diagnostic {
-            severity: higher_severity(error.severity, next_error.severity),
-            message: formatdoc!(
-                "
-                    Both options failed
-                    First variant: {}
-                    Second variant: {}
-                    Resolve at least one of them
-                ",
-                error.message,
-                next_error.message
-            ),
-        })
-    })
+    formatdoc! {"
+        All options failed
+        {}
+        Resolve at least one of them
+    ", formatted_variants.join("\n")}
+}
+
+/// The `branch` macro is used to evaluate multiple expressions and return the first successful result.
+/// If all expressions fail, it collects the error messages and returns a combined error.
+///
+/// This macro is used instead of a function because it can perform lazy evaluation and has better readability.
+#[macro_export]
+macro_rules! branch {
+    ($($result:expr),+) => {{
+        let mut messages = Vec::new();
+        let mut result = None;
+
+        $(
+            if result.is_none() {
+                match $result {
+                    Ok(val) => {
+                        result = Some(val);
+                    },
+                    Err(err) => {
+                        messages.push(err);
+                    },
+                }
+            }
+        )+
+
+        if let Some(result) = result {
+            Ok(result)
+        } else {
+            Err(Diagnostic {
+                message: $crate::utils::format_error_message(&messages),
+                severity: messages.into_iter().fold(Severity::Warning, |acc, diagnostic| $crate::utils::higher_severity(acc, diagnostic.severity))
+            })
+        }
+    }};
 }