Add --partition flag

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 {
@@ -149,7 +151,7 @@ pub struct TestArgs {
     run_native: bool,
 
     /// 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.
@@ -213,6 +215,11 @@ pub struct TestArgs {
     #[arg(long, value_enum, default_value_t)]
     tracked_resource: ForgeTrackedResource,
 
+    /// If specified, divides tests into partitions and runs specified partition.
+    /// <PARTITION> is in the format INDEX/TOTAL, where INDEX is the 1-based index of the partition to run, and TOTAL is the number of partitions.
+    #[arg(long, conflicts_with = "exact")]
+    partition: Option<Partition>,
+
     /// Additional arguments for cairo-coverage or cairo-profiler
     #[arg(last = true)]
     additional_args: Vec<OsString>,
@@ -327,7 +334,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,82 @@
+use serde::Serialize;
+use std::{collections::HashMap, str::FromStr};
+
+#[derive(Debug, Clone, Copy, Serialize)]
+#[non_exhaustive]
+pub struct Partition {
+    pub index: usize,
+    pub total: usize,
+}
+
+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 })
+    }
+}
+
+/// A mapping between test full paths and their assigned partition indices.
+#[derive(Serialize)]
+pub struct TestPartitionMap(HashMap<String, usize>);
+
+#[derive(Serialize)]
+#[non_exhaustive]
+pub struct PartitionConfig {
+    pub partition: Partition,
+    pub test_partition_map: TestPartitionMap,
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use test_case::test_case;
+
+    #[test]
+    fn test_happy_case() {
+        let partition = "2/5".parse::<Partition>().unwrap();
+        assert_eq!(partition.index, 2);
+        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]
+    #[test_case("-1/5", "INDEX")]
+    #[test_case("2/-5", "TOTAL")]
+    #[test_case("a/5", "INDEX")]
+    #[test_case("2/b", "TOTAL")]
+    fn test_invalid_integer(partition: &str, invalid_part: &str) {
+        let err = partition.parse::<Partition>().unwrap_err();
+        assert_eq!(err, format!("{invalid_part} must be a positive integer"));
+    }
+
+    #[test_case("0/5")]
+    #[test_case("6/5")]
+    #[test_case("2/0")]
+    fn test_out_of_bounds(partition: &str) {
+        let err = partition.parse::<Partition>().unwrap_err();
+        assert_eq!(err, "Invalid partition values: ensure 1 <= INDEX <= TOTAL");
+    }
+}

crates/forge/tests/e2e/mod.rs

@@ -24,6 +24,7 @@ mod fuzzing;
 mod io_operations;
 mod new;
 mod package_warnings;
+mod partitioning;
 mod plugin_diagnostics;
 mod plugin_versions;
 mod profiles;

crates/forge/tests/e2e/partitioning.rs

@@ -0,0 +1,19 @@
+use crate::e2e::common::runner::{setup_package, test_runner};
+use indoc::indoc;
+use shared::test_utils::output_assert::assert_stderr_contains;
+
+#[test]
+fn test_does_not_work_with_exact_flag() {
+    let temp = setup_package("simple_package");
+    let output = test_runner(&temp)
+        .args(["--partition", "3/3", "--workspace", "--exact"])
+        .assert()
+        .code(2);
+
+    assert_stderr_contains(
+        output,
+        indoc! {r"
+        error: the argument '--partition <PARTITION>' cannot be used with '--exact'
+    "},
+    );
+}