Add partitioning logic

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

@@ -128,6 +128,8 @@ pub enum TestCaseSummary<T: TestType> {
     },
     /// Test case skipped due to exit first or execution interrupted, test result is ignored.
     Interrupted {},
+    /// Test case excluded from current partition
+    ExcludedFromPartition {},
 }
 
 #[derive(Debug)]
@@ -145,7 +147,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
+            }
         }
     }
 
@@ -227,6 +231,7 @@ impl TestCaseSummary<Fuzzing> {
             },
             TestCaseSummary::Ignored { name } => TestCaseSummary::Ignored { name: name.clone() },
             TestCaseSummary::Interrupted {} => TestCaseSummary::Interrupted {},
+            TestCaseSummary::ExcludedFromPartition {} => TestCaseSummary::ExcludedFromPartition {},
         }
     }
 }
@@ -454,6 +459,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/partition.rs

@@ -1,13 +1,16 @@
+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};
 
+use crate::run_tests::package::RunForPackageArgs;
+
 #[derive(Debug, Clone, Copy, Serialize)]
 pub struct Partition {
     index: usize,
     total: usize,
 }
 
-#[allow(dead_code)] // TODO: Removed in later PRs
 impl Partition {
     #[must_use]
     pub fn index(&self) -> usize {
@@ -49,15 +52,57 @@ impl FromStr for Partition {
 #[derive(Serialize)]
 pub struct TestPartitionMap(HashMap<String, usize>);
 
+impl TestPartitionMap {
+    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,
     test_partition_map: TestPartitionMap,
 }
 
-#[allow(dead_code)] // TODO: Removed in later PRs
 impl PartitionConfig {
-    pub fn new(partition: Partition, test_partition_map: TestPartitionMap) -> Self {
+    pub fn new(partition: Partition, packages_args: &[RunForPackageArgs]) -> Self {
+        let test_partition_map = TestPartitionMap::from_packages_args(packages_args, partition);
         Self {
             partition,
             test_partition_map,

crates/forge/src/run_tests/package.rs

@@ -6,6 +6,7 @@ use crate::{
     TestArgs,
     block_number_map::BlockNumberMap,
     combine_configs::combine_configs,
+    partition::PartitionConfig,
     run_tests::messages::{
         collected_tests_count::CollectedTestsCountMessage, tests_run::TestsRunMessage,
         tests_summary::TestsSummaryMessage,
@@ -28,7 +29,11 @@ use configuration::load_package_config;
 use console::Style;
 use forge_runner::{
     forge_config::ForgeConfig,
-    package_tests::{raw::TestTargetRaw, with_config_resolved::TestTargetWithResolvedConfig},
+    package_tests::{
+        TestCase, TestTarget,
+        raw::TestTargetRaw,
+        with_config_resolved::{TestCaseResolvedConfig, TestTargetWithResolvedConfig},
+    },
     running::with_config::test_target_with_config,
     test_case_summary::AnyTestCaseSummary,
     test_target_summary::TestTargetSummary,
@@ -163,8 +168,34 @@ async fn test_package_with_config_resolved(
     Ok(test_targets_with_resolved_config)
 }
 
-fn sum_test_cases(test_targets: &[TestTargetWithResolvedConfig]) -> usize {
-    test_targets.iter().map(|tc| tc.test_cases.len()).sum()
+fn sum_test_cases_from_package(
+    test_targets: &[TestTarget<TestCaseResolvedConfig>],
+    partition_config: Option<&PartitionConfig>,
+) -> usize {
+    test_targets
+        .iter()
+        .map(|tt| sum_test_cases_from_test_target(tt.test_cases.clone(), partition_config))
+        .sum()
+}
+
+fn sum_test_cases_from_test_target(
+    test_cases: Vec<TestCase<TestCaseResolvedConfig>>,
+    partition_config: Option<&PartitionConfig>,
+) -> usize {
+    if let Some(partition_config) = partition_config {
+        test_cases
+            .into_iter()
+            .filter(|test_case| {
+                partition_config.partition().index()
+                    == *partition_config
+                        .test_partition_map()
+                        .get(&test_case.name)
+                        .expect("Test case name not found in partitions mapping")
+            })
+            .count()
+    } else {
+        test_cases.len()
+    }
 }
 
 #[tracing::instrument(skip_all, level = "debug")]
@@ -177,6 +208,7 @@ pub async fn run_for_package(
         package_name,
     }: RunForPackageArgs,
     block_number_map: &mut BlockNumberMap,
+    partition_config: Option<&PartitionConfig>,
     ui: Arc<UI>,
 ) -> Result<PackageTestResult> {
     let mut test_targets = test_package_with_config_resolved(
@@ -187,7 +219,7 @@ pub async fn run_for_package(
         &tests_filter,
     )
     .await?;
-    let all_tests = sum_test_cases(&test_targets);
+    let all_tests = sum_test_cases_from_package(&test_targets, partition_config);
 
     for test_target in &mut test_targets {
         tests_filter.filter_tests(&mut test_target.test_cases)?;
@@ -196,7 +228,8 @@ pub async fn run_for_package(
     warn_if_available_gas_used_with_incompatible_scarb_version(&test_targets, &ui)?;
     warn_if_incompatible_rpc_version(&test_targets, ui.clone()).await?;
 
-    let not_filtered = sum_test_cases(&test_targets);
+    let not_filtered = sum_test_cases_from_package(&test_targets, partition_config);
+
     ui.println(&CollectedTestsCountMessage {
         tests_num: not_filtered,
         package_name: package_name.clone(),
@@ -208,11 +241,17 @@ pub async fn run_for_package(
         let ui = ui.clone();
         ui.println(&TestsRunMessage::new(
             test_target.tests_location,
-            test_target.test_cases.len(),
+            sum_test_cases_from_test_target(test_target.test_cases.clone(), partition_config),
         ));
 
-        let summary =
-            run_for_test_target(test_target, forge_config.clone(), &tests_filter, ui).await?;
+        let summary = run_for_test_target(
+            test_target,
+            forge_config.clone(),
+            &tests_filter,
+            partition_config,
+            ui,
+        )
+        .await?;
 
         match summary {
             TestTargetRunResult::Ok(summary) => {

crates/forge/src/run_tests/test_target.rs

@@ -1,3 +1,4 @@
+use crate::partition::PartitionConfig;
 use anyhow::Result;
 use forge_runner::messages::TestResultMessage;
 use forge_runner::{
@@ -25,6 +26,7 @@ pub async fn run_for_test_target(
     tests: TestTargetWithResolvedConfig,
     forge_config: Arc<ForgeConfig>,
     tests_filter: &impl TestCaseFilter,
+    partition_config: Option<&PartitionConfig>,
     ui: Arc<UI>,
 ) -> Result<TestTargetRunResult> {
     let casm_program = tests.casm_program.clone();
@@ -40,6 +42,28 @@ pub async fn run_for_test_target(
     for case in tests.test_cases {
         let case_name = case.name.clone();
 
+        if let Some(partition_config) = &partition_config {
+            let test_partition = partition_config
+                .test_partition_map()
+                .get(&case.name)
+                .unwrap_or_else(|| {
+                    unreachable!(
+                        "Test '{}' should have been mapped to a partition",
+                        case.name
+                    )
+                });
+            let is_included_in_partition = *test_partition == partition_config.partition().index();
+
+            if !is_included_in_partition {
+                tasks.push(tokio::task::spawn(async {
+                    Ok(AnyTestCaseSummary::Single(
+                        TestCaseSummary::ExcludedFromPartition {},
+                    ))
+                }));
+                continue;
+            }
+        }
+
         if !tests_filter.should_be_run(&case) {
             tasks.push(tokio::task::spawn(async {
                 // TODO TestCaseType should also be encoded in the test case definition
@@ -68,7 +92,7 @@ pub async fn run_for_test_target(
     while let Some(task) = tasks.next().await {
         let result = task??;
 
-        if !result.is_interrupted() {
+        if !result.is_interrupted() && !result.is_excluded_from_partition() {
             let test_result_message = TestResultMessage::new(
                 &result,
                 forge_config.output_config.detailed_resources,