Tests partitioning

crates/forge-runner/src/messages.rs

@@ -13,6 +13,7 @@ enum TestResultStatus {
     Failed,
     Ignored,
     Interrupted,
+    ExcludedFromPartition,
 }
 
 impl From<&AnyTestCaseSummary> for TestResultStatus {
@@ -26,6 +27,10 @@ impl From<&AnyTestCaseSummary> for TestResultStatus {
             | AnyTestCaseSummary::Fuzzing(TestCaseSummary::Ignored { .. }) => Self::Ignored,
             AnyTestCaseSummary::Single(TestCaseSummary::Interrupted { .. })
             | AnyTestCaseSummary::Fuzzing(TestCaseSummary::Interrupted { .. }) => Self::Interrupted,
+            AnyTestCaseSummary::Single(TestCaseSummary::ExcludedFromPartition { .. })
+            | AnyTestCaseSummary::Fuzzing(TestCaseSummary::ExcludedFromPartition { .. }) => {
+                Self::ExcludedFromPartition
+            }
         }
     }
 }
@@ -110,7 +115,9 @@ impl TestResultMessage {
             match self.status {
                 TestResultStatus::Passed => return format!("\n\n{msg}"),
                 TestResultStatus::Failed => return format!("\n\nFailure data:{msg}"),
-                TestResultStatus::Ignored | TestResultStatus::Interrupted => return String::new(),
+                TestResultStatus::Ignored
+                | TestResultStatus::Interrupted
+                | TestResultStatus::ExcludedFromPartition => return String::new(),
             }
         }
         String::new()
@@ -124,6 +131,11 @@ impl TestResultMessage {
             TestResultStatus::Interrupted => {
                 unreachable!("Interrupted tests should not have visible message representation")
             }
+            TestResultStatus::ExcludedFromPartition => {
+                unreachable!(
+                    "Tests excluded from partition should not have visible message representation"
+                )
+            }
         }
     }
 }

crates/forge-runner/src/package_tests/with_config_resolved.rs

@@ -13,7 +13,8 @@ pub type TestTargetWithResolvedConfig = TestTarget<TestCaseResolvedConfig>;
 
 pub type TestCaseWithResolvedConfig = TestCase<TestCaseResolvedConfig>;
 
-fn sanitize_test_case_name(name: &str) -> String {
+#[must_use]
+pub fn sanitize_test_case_name(name: &str) -> String {
     // Test names generated by `#[test]` and `#[fuzzer]` macros contain internal suffixes
     name.replace("__snforge_internal_test_generated", "")
         .replace("__snforge_internal_fuzzer_generated", "")

crates/forge-runner/src/test_case_summary.rs

@@ -163,6 +163,8 @@ pub enum TestCaseSummary<T: TestType> {
         /// Name of the test case
         name: String,
     },
+    /// Test case excluded from current partition
+    ExcludedFromPartition {},
     /// Test case skipped due to exit first or execution interrupted, test result is ignored.
     Interrupted {},
 }
@@ -182,7 +184,9 @@ impl<T: TestType> TestCaseSummary<T> {
             TestCaseSummary::Failed { name, .. }
             | TestCaseSummary::Passed { name, .. }
             | TestCaseSummary::Ignored { name, .. } => Some(name),
-            TestCaseSummary::Interrupted { .. } => None,
+            TestCaseSummary::Interrupted { .. } | TestCaseSummary::ExcludedFromPartition { .. } => {
+                None
+            }
         }
     }
 
@@ -264,6 +268,7 @@ impl TestCaseSummary<Fuzzing> {
             },
             TestCaseSummary::Ignored { name } => TestCaseSummary::Ignored { name: name.clone() },
             TestCaseSummary::Interrupted {} => TestCaseSummary::Interrupted {},
+            TestCaseSummary::ExcludedFromPartition {} => TestCaseSummary::ExcludedFromPartition {},
         }
     }
 }
@@ -491,6 +496,15 @@ impl AnyTestCaseSummary {
                 | AnyTestCaseSummary::Fuzzing(TestCaseSummary::Ignored { .. })
         )
     }
+
+    #[must_use]
+    pub fn is_excluded_from_partition(&self) -> bool {
+        matches!(
+            self,
+            AnyTestCaseSummary::Single(TestCaseSummary::ExcludedFromPartition { .. })
+                | AnyTestCaseSummary::Fuzzing(TestCaseSummary::ExcludedFromPartition { .. })
+        )
+    }
 }
 
 #[cfg(test)]

crates/forge/src/lib.rs

@@ -1,4 +1,5 @@
 use crate::compatibility_check::{Requirement, RequirementsChecker, create_version_parser};
+use crate::partition::Partition;
 use anyhow::Result;
 use camino::Utf8PathBuf;
 use clap::{CommandFactory, Parser, Subcommand, ValueEnum};
@@ -25,6 +26,7 @@ mod clean;
 mod combine_configs;
 mod compatibility_check;
 mod new;
+mod partition;
 mod profile_validation;
 pub mod run_tests;
 pub mod scarb;
@@ -83,7 +85,7 @@ enum ForgeSubcommand {
     /// Run tests for a project in the current directory
     Test {
         #[command(flatten)]
-        args: TestArgs,
+        args: Box<TestArgs>,
     },
     /// Create a new Forge project at <PATH>
     New {
@@ -142,7 +144,7 @@ pub struct TestArgs {
     trace_args: TraceArgs,
 
     /// Use exact matches for `test_filter`
-    #[arg(short, long)]
+    #[arg(short, long, conflicts_with = "partition")]
     exact: bool,
 
     /// Skips any tests whose name contains the given SKIP string.
@@ -203,6 +205,10 @@ pub struct TestArgs {
     #[arg(long, value_enum, default_value_t)]
     tracked_resource: ForgeTrackedResource,
 
+    /// If specified, divides tests into `total` partitions and runs only the partition with the given `index` (1-based).
+    #[arg(long, conflicts_with = "exact")]
+    partition: Option<Partition>,
+
     /// Additional arguments for cairo-coverage or cairo-profiler
     #[arg(last = true)]
     additional_args: Vec<OsString>,
@@ -297,7 +303,7 @@ pub fn main_execution(ui: Arc<UI>) -> Result<ExitStatus> {
                 .enable_all()
                 .build()?;
 
-            rt.block_on(run_for_workspace(args, ui))
+            rt.block_on(run_for_workspace(*args, ui))
         }
         ForgeSubcommand::CheckRequirements => {
             check_requirements(true, &ui)?;

crates/forge/src/partition.rs

@@ -0,0 +1,162 @@
+use crate::run_tests::package::RunForPackageArgs;
+use cairo_lang_sierra::ids::FunctionId;
+use forge_runner::package_tests::with_config_resolved::sanitize_test_case_name;
+use serde::Serialize;
+use std::{collections::HashMap, str::FromStr};
+
+#[derive(Debug, Clone, Copy, Serialize)]
+pub struct Partition {
+    index: usize,
+    total: usize,
+}
+
+impl Partition {
+    #[must_use]
+    pub fn index_0_based(&self) -> usize {
+        self.index - 1
+    }
+
+    #[must_use]
+    pub fn index_1_based(&self) -> usize {
+        self.index
+    }
+
+    #[must_use]
+    pub fn total(&self) -> usize {
+        self.total
+    }
+}
+
+impl FromStr for Partition {
+    type Err = String;
+
+    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
+        let parts: Vec<&str> = s.split('/').collect();
+        if parts.len() != 2 {
+            return Err("Partition must be in the format <INDEX>/<TOTAL>".to_string());
+        }
+
+        let index = parts[0]
+            .parse::<usize>()
+            .map_err(|_| "INDEX must be a positive integer".to_string())?;
+        let total = parts[1]
+            .parse::<usize>()
+            .map_err(|_| "TOTAL must be a positive integer".to_string())?;
+
+        if index == 0 || total == 0 || index > total {
+            return Err("Invalid partition values: ensure 1 <= INDEX <= TOTAL".to_string());
+        }
+
+        Ok(Partition { index, total })
+    }
+}
+
+#[derive(Serialize)]
+pub struct TestsPartitionsMapping(HashMap<String, usize>);
+
+impl TestsPartitionsMapping {
+    pub fn get(&self, test_name: &str) -> Option<&usize> {
+        self.0.get(test_name)
+    }
+
+    pub fn insert(&mut self, test_name: String, partition_index: usize) {
+        self.0.insert(test_name, partition_index);
+    }
+
+    pub fn from_packages_args(packages_args: &[RunForPackageArgs], partition: Partition) -> Self {
+        let mut full_paths: Vec<String> = packages_args
+            .iter()
+            .flat_map(|pkg| pkg.test_targets.iter())
+            .flat_map(|tt| {
+                tt.sierra_program
+                    .debug_info
+                    .as_ref()
+                    .and_then(|info| info.executables.get("snforge_internal_test_executable"))
+                    .into_iter()
+                    .flatten()
+            })
+            .filter_map(|fid: &FunctionId| {
+                fid.debug_name
+                    .as_ref()
+                    .map(std::string::ToString::to_string)
+            })
+            .collect();
+
+        full_paths.sort();
+
+        let total = partition.total();
+        let mut mapping = HashMap::with_capacity(full_paths.len());
+
+        for (i, path) in full_paths.into_iter().enumerate() {
+            let partition_index_1_based = (i % total) + 1;
+            mapping.insert(sanitize_test_case_name(&path), partition_index_1_based);
+        }
+
+        Self(mapping)
+    }
+}
+
+#[derive(Serialize)]
+pub struct PartitionConfig {
+    partition: Partition,
+    partitions_mapping: TestsPartitionsMapping,
+}
+
+impl PartitionConfig {
+    pub fn new(partition: Partition, packages_args: &[RunForPackageArgs]) -> Self {
+        let partitions_mapping =
+            TestsPartitionsMapping::from_packages_args(packages_args, partition);
+        Self {
+            partition,
+            partitions_mapping,
+        }
+    }
+
+    pub fn partition(&self) -> Partition {
+        self.partition
+    }
+
+    pub fn partitions_mapping(&self) -> &TestsPartitionsMapping {
+        &self.partitions_mapping
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_happy_case() {
+        let partition = "2/5".parse::<Partition>().unwrap();
+        assert_eq!(partition.index_1_based(), 2);
+        assert_eq!(partition.index_0_based(), 1);
+        assert_eq!(partition.total(), 5);
+    }
+
+    #[test]
+    fn test_invalid_format() {
+        let err = "2-5".parse::<Partition>().unwrap_err();
+        assert_eq!(err, "Partition must be in the format <INDEX>/<TOTAL>");
+    }
+
+    #[test]
+    fn test_non_integer() {
+        let err = "a/5".parse::<Partition>().unwrap_err();
+        assert_eq!(err, "INDEX must be a positive integer");
+
+        let err = "2/b".parse::<Partition>().unwrap_err();
+        assert_eq!(err, "TOTAL must be a positive integer");
+    }
+
+    #[test]
+    fn test_out_of_bounds() {
+        let err = "0/5".parse::<Partition>().unwrap_err();
+        assert_eq!(err, "Invalid partition values: ensure 1 <= INDEX <= TOTAL");
+
+        let err = "6/5".parse::<Partition>().unwrap_err();
+        assert_eq!(err, "Invalid partition values: ensure 1 <= INDEX <= TOTAL");
+
+        let err = "2/0".parse::<Partition>().unwrap_err();
+        assert_eq!(err, "Invalid partition values: ensure 1 <= INDEX <= TOTAL");
+    }
+}

crates/forge/src/run_tests/messages/mod.rs

@@ -1,6 +1,7 @@
 pub mod collected_tests_count;
 pub mod latest_blocks_numbers;
 pub mod overall_summary;
+pub mod partition;
 pub mod tests_failure_summary;
 pub mod tests_run;
 pub mod tests_summary;

crates/forge/src/run_tests/messages/partition.rs

@@ -0,0 +1,33 @@
+use foundry_ui::Message;
+use serde::Serialize;
+use serde_json::{Value, json};
+
+use crate::partition::Partition;
+
+#[derive(Serialize)]
+pub struct PartitionMessage {
+    partition: Partition,
+}
+
+impl PartitionMessage {
+    #[must_use]
+    pub fn new(partition: Partition) -> Self {
+        Self { partition }
+    }
+}
+
+impl Message for PartitionMessage {
+    fn text(&self) -> String {
+        format!(
+            "Finished partition run: {}/{}",
+            self.partition.index_1_based(),
+            self.partition.total()
+        )
+    }
+
+    fn json(&self) -> Value {
+        json!({
+            "partition": format!("{}/{}", self.partition.index_1_based(), self.partition.total())
+        })
+    }
+}