Pass args to coverage and profiler

CHANGELOG.md

@@ -7,6 +7,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Forge
+
+#### Changed
+
+- You can now pass arguments to `cairo-profiler` or `cairo-coverage`,
+  e.g. `snforge test --coverage --include tests-functions`. If flags name duplicates a flag in `snforge test`, you can
+  use `--` to separate them, e.g. `snforge test --coverage -- --help`
+- You can't use now `--coverage` and `--build-profile` flags at the same time. If you want to use both, you need to run
+  `snforge test` twice with different flags.
+
 ## [0.32.0] - 2024-10-16
 
 ### Cast

crates/forge-runner/src/coverage_api.rs

@@ -3,6 +3,7 @@ use indoc::{formatdoc, indoc};
 use scarb_api::metadata::Metadata;
 use semver::Version;
 use shared::command::CommandExt;
+use std::ffi::OsString;
 use std::process::Stdio;
 use std::{env, fs, path::PathBuf, process::Command};
 use toml_edit::{DocumentMut, Table};
@@ -19,7 +20,7 @@ const CAIRO_COVERAGE_REQUIRED_ENTRIES: [(&str, &str); 3] = [
     ("inlining-strategy", "\"avoid\""),
 ];
 
-pub fn run_coverage(saved_trace_data_paths: &[PathBuf]) -> Result<()> {
+pub fn run_coverage(saved_trace_data_paths: &[PathBuf], coverage_args: &[OsString]) -> Result<()> {
     let coverage = env::var("CAIRO_COVERAGE")
         .map(PathBuf::from)
         .ok()
@@ -34,10 +35,6 @@ pub fn run_coverage(saved_trace_data_paths: &[PathBuf]) -> Result<()> {
         }
     );
 
-    let dir_to_save_coverage = PathBuf::from(COVERAGE_DIR);
-    fs::create_dir_all(&dir_to_save_coverage).context("Failed to create a coverage dir")?;
-    let path_to_save_coverage = dir_to_save_coverage.join(OUTPUT_FILE_NAME);
-
     let trace_files: Vec<&str> = saved_trace_data_paths
         .iter()
         .map(|trace_data_path| {
@@ -47,10 +44,19 @@ pub fn run_coverage(saved_trace_data_paths: &[PathBuf]) -> Result<()> {
         })
         .collect();
 
-    Command::new(coverage)
-        .arg("--output-path")
-        .arg(&path_to_save_coverage)
+    let mut command = Command::new(coverage);
+
+    if coverage_args.iter().all(|arg| arg != "--output-path") {
+        let dir_to_save_coverage = PathBuf::from(COVERAGE_DIR);
+        fs::create_dir_all(&dir_to_save_coverage).context("Failed to create a coverage dir")?;
+        let path_to_save_coverage = dir_to_save_coverage.join(OUTPUT_FILE_NAME);
+
+        command.arg("--output-path").arg(&path_to_save_coverage);
+    }
+
+    command
         .args(trace_files)
+        .args(coverage_args)
         .stdout(Stdio::inherit())
         .stderr(Stdio::inherit())
         .output_checked()

crates/forge-runner/src/forge_config.rs

@@ -1,6 +1,7 @@
 use camino::Utf8PathBuf;
 use cheatnet::runtime_extensions::forge_runtime_extension::contracts_data::ContractsData;
 use std::collections::HashMap;
+use std::ffi::OsString;
 use std::num::NonZeroU32;
 use std::sync::Arc;
 
@@ -29,20 +30,27 @@ pub struct OutputConfig {
     pub versioned_programs_dir: Utf8PathBuf,
 }
 
-#[derive(Debug, PartialEq, Clone, Copy, Default)]
+#[derive(Debug, PartialEq, Clone, Default)]
 pub struct ExecutionDataToSave {
     pub trace: bool,
     pub profile: bool,
     pub coverage: bool,
+    pub additional_args: Vec<OsString>,
 }
 
 impl ExecutionDataToSave {
     #[must_use]
-    pub fn from_flags(save_trace_data: bool, build_profile: bool, coverage: bool) -> Self {
+    pub fn from_flags(
+        save_trace_data: bool,
+        build_profile: bool,
+        coverage: bool,
+        additional_args: &[OsString],
+    ) -> Self {
         Self {
             trace: save_trace_data,
             profile: build_profile,
             coverage,
+            additional_args: additional_args.to_vec(),
         }
     }
     #[must_use]

crates/forge-runner/src/lib.rs

@@ -60,7 +60,7 @@ pub trait TestCaseFilter {
 
 pub fn maybe_save_trace_and_profile(
     result: &AnyTestCaseSummary,
-    execution_data_to_save: ExecutionDataToSave,
+    execution_data_to_save: &ExecutionDataToSave,
 ) -> Result<Option<PathBuf>> {
     if let AnyTestCaseSummary::Single(TestCaseSummary::Passed {
         name, trace_data, ..
@@ -69,7 +69,7 @@ pub fn maybe_save_trace_and_profile(
         if execution_data_to_save.is_vm_trace_needed() {
             let trace_path = save_trace_data(name, trace_data)?;
             if execution_data_to_save.profile {
-                run_profiler(name, &trace_path)?;
+                run_profiler(name, &trace_path, &execution_data_to_save.additional_args)?;
             }
             return Ok(Some(trace_path));
         }
@@ -78,21 +78,24 @@ pub fn maybe_save_trace_and_profile(
 }
 
 pub fn maybe_generate_coverage(
-    execution_data_to_save: ExecutionDataToSave,
+    execution_data_to_save: &ExecutionDataToSave,
     saved_trace_data_paths: &[PathBuf],
 ) -> Result<()> {
     if execution_data_to_save.coverage {
         if saved_trace_data_paths.is_empty() {
             print_as_warning(&anyhow!("No trace data to generate coverage from"));
         } else {
-            run_coverage(saved_trace_data_paths)?;
+            run_coverage(
+                saved_trace_data_paths,
+                &execution_data_to_save.additional_args,
+            )?;
         }
     }
     Ok(())
 }
 
 pub fn maybe_save_versioned_program(
-    execution_data_to_save: ExecutionDataToSave,
+    execution_data_to_save: &ExecutionDataToSave,
     test_target: &TestTargetWithResolvedConfig,
     versioned_programs_dir: &Utf8Path,
     package_name: &str,

crates/forge-runner/src/profiler_api.rs

@@ -1,23 +1,34 @@
 use anyhow::{Context, Result};
 use shared::command::CommandExt;
+use std::ffi::OsString;
 use std::process::Stdio;
 use std::{env, fs, path::PathBuf, process::Command};
 
 pub const PROFILE_DIR: &str = "profile";
 
-pub fn run_profiler(test_name: &str, trace_path: &PathBuf) -> Result<()> {
+pub fn run_profiler(
+    test_name: &str,
+    trace_path: &PathBuf,
+    profiler_args: &[OsString],
+) -> Result<()> {
     let profiler = env::var("CAIRO_PROFILER")
         .map(PathBuf::from)
         .ok()
         .unwrap_or_else(|| PathBuf::from("cairo-profiler"));
-    let dir_to_save_profile = PathBuf::from(PROFILE_DIR);
-    fs::create_dir_all(&dir_to_save_profile).context("Failed to create a profile dir")?;
-    let path_to_save_profile = dir_to_save_profile.join(format!("{test_name}.pb.gz"));
 
-    Command::new(profiler)
+    let mut command = Command::new(profiler);
+
+    if profiler_args.iter().all(|arg| arg != "--output-path") {
+        let dir_to_save_profile = PathBuf::from(PROFILE_DIR);
+        fs::create_dir_all(&dir_to_save_profile).context("Failed to create a profile dir")?;
+        let path_to_save_profile = dir_to_save_profile.join(format!("{test_name}.pb.gz"));
+
+        command.arg("--output-path").arg(&path_to_save_profile);
+    }
+
+    command
         .arg(trace_path)
-        .arg("--output-path")
-        .arg(&path_to_save_profile)
+        .args(profiler_args)
         .stdout(Stdio::inherit())
         .stderr(Stdio::inherit())
         .output_checked()

crates/forge/src/combine_configs.rs

@@ -6,6 +6,7 @@ use forge_runner::forge_config::{
 };
 use rand::{thread_rng, RngCore};
 use std::env;
+use std::ffi::OsString;
 use std::num::NonZeroU32;
 use std::sync::Arc;
 
@@ -24,11 +25,13 @@ pub fn combine_configs(
     cache_dir: Utf8PathBuf,
     versioned_programs_dir: Utf8PathBuf,
     forge_config_from_scarb: &ForgeConfigFromScarb,
+    coverage_or_profiler_args: &[OsString],
 ) -> ForgeConfig {
     let execution_data_to_save = ExecutionDataToSave::from_flags(
         save_trace_data || forge_config_from_scarb.save_trace_data,
         build_profile || forge_config_from_scarb.build_profile,
         coverage || forge_config_from_scarb.coverage,
+        coverage_or_profiler_args,
     );
 
     ForgeConfig {
@@ -73,6 +76,7 @@ mod tests {
             Default::default(),
             Default::default(),
             &Default::default(),
+            &[],
         );
         let config2 = combine_configs(
             false,
@@ -87,6 +91,7 @@ mod tests {
             Default::default(),
             Default::default(),
             &Default::default(),
+            &[],
         );
 
         assert_ne!(config.test_runner_config.fuzzer_seed, 0);
@@ -112,6 +117,7 @@ mod tests {
             Default::default(),
             Default::default(),
             &Default::default(),
+            &[],
         );
         assert_eq!(
             config,
@@ -162,6 +168,7 @@ mod tests {
             Default::default(),
             Default::default(),
             &config_from_scarb,
+            &[],
         );
         assert_eq!(
             config,
@@ -182,6 +189,7 @@ mod tests {
                         trace: true,
                         profile: true,
                         coverage: true,
+                        additional_args: vec![],
                     },
                     versioned_programs_dir: Default::default(),
                 }),
@@ -215,6 +223,7 @@ mod tests {
             Default::default(),
             Default::default(),
             &config_from_scarb,
+            &[],
         );
 
         assert_eq!(
@@ -236,6 +245,7 @@ mod tests {
                         trace: true,
                         profile: true,
                         coverage: true,
+                        additional_args: vec![],
                     },
                     versioned_programs_dir: Default::default(),
                 }),

crates/forge/src/lib.rs

@@ -4,6 +4,7 @@ use forge_runner::CACHE_DIR;
 use run_tests::workspace::run_for_workspace;
 use scarb_api::{metadata::MetadataCommandExt, ScarbCommand};
 use scarb_ui::args::{FeaturesSpec, PackagesFilter};
+use std::ffi::OsString;
 use std::{fs, num::NonZeroU32, thread::available_parallelism};
 use tokio::runtime::Builder;
 use universal_sierra_compiler_api::UniversalSierraCompilerCommand;
@@ -84,6 +85,7 @@ enum ColorOption {
 #[allow(clippy::struct_excessive_bools)]
 pub struct TestArgs {
     /// Name used to filter tests
+    #[clap(allow_hyphen_values = true)]
     test_filter: Option<String>,
     /// Use exact matches for `test_filter`
     #[arg(short, long)]
@@ -127,11 +129,11 @@ pub struct TestArgs {
     save_trace_data: bool,
 
     /// Build profiles of all tests which have passed and are not fuzz tests using the cairo-profiler
-    #[arg(long)]
+    #[arg(long, conflicts_with = "coverage")]
     build_profile: bool,
 
     /// Generate a coverage report for the executed tests which have passed and are not fuzz tests using the cairo-coverage
-    #[arg(long)]
+    #[arg(long, conflicts_with = "build_profile")]
     coverage: bool,
 
     /// Number of maximum steps during a single test. For fuzz tests this value is applied to each subtest separately.
@@ -145,6 +147,10 @@ pub struct TestArgs {
     /// Build contracts separately in the scarb starknet contract target
     #[arg(long)]
     no_optimization: bool,
+
+    /// Additional arguments for coverage or profiler
+    #[clap(allow_hyphen_values = true)]
+    coverage_or_profiler_args: Vec<OsString>,
 }
 
 pub enum ExitStatus {
@@ -173,7 +179,16 @@ pub fn main_execution() -> Result<ExitStatus> {
 
             Ok(ExitStatus::Success)
         }
-        ForgeSubcommand::Test { args } => {
+        ForgeSubcommand::Test { mut args } => {
+            // clap can capture the first flag in test filter,
+            // let's move it to coverage_or_profiler_args
+            if let Some(test_filter) = &args.test_filter {
+                if test_filter.starts_with('-') {
+                    args.coverage_or_profiler_args.insert(0, test_filter.into());
+                    args.test_filter = None;
+                }
+            }
+
             let cores = if let Ok(available_cores) = available_parallelism() {
                 available_cores.get()
             } else {

crates/forge/src/run_tests/package.rs

@@ -78,6 +78,7 @@ impl RunForPackageArgs {
             cache_dir.clone(),
             versioned_programs_dir,
             &forge_config_from_scarb,
+            &args.coverage_or_profiler_args,
         ));
 
         let test_filter = TestsFilter::from_flags(

crates/forge/src/run_tests/test_target.rs

@@ -39,7 +39,7 @@ pub async fn run_for_test_target(
     let (send, mut rec) = channel(1);
 
     let maybe_versioned_program_path = Arc::new(maybe_save_versioned_program(
-        forge_config.output_config.execution_data_to_save,
+        &forge_config.output_config.execution_data_to_save,
         &tests,
         &forge_config.output_config.versioned_programs_dir,
         package_name,
@@ -95,7 +95,7 @@ pub async fn run_for_test_target(
 
         let trace_path = maybe_save_trace_and_profile(
             &result,
-            forge_config.output_config.execution_data_to_save,
+            &forge_config.output_config.execution_data_to_save,
         )?;
         if let Some(path) = trace_path {
             saved_trace_data_paths.push(path);
@@ -110,7 +110,7 @@ pub async fn run_for_test_target(
     }
 
     maybe_generate_coverage(
-        forge_config.output_config.execution_data_to_save,
+        &forge_config.output_config.execution_data_to_save,
         &saved_trace_data_paths,
     )?;
 

crates/forge/tests/e2e/build_profile.rs

@@ -27,3 +27,17 @@ fn simple_package_build_profile() {
     // Check if it doesn't crash in case some data already exists
     test_runner(&temp).arg("--build-profile").assert().code(1);
 }
+
+#[test]
+fn simple_package_build_profile_and_pass_args() {
+    let temp = setup_package("simple_package");
+
+    test_runner(&temp)
+        .arg("--build-profile")
+        .arg("--output-path")
+        .arg("my_file.pb.gz")
+        .assert()
+        .code(1);
+
+    assert!(temp.join("my_file.pb.gz").is_file());
+}