diff --git a/crates/forge/src/lib.rs b/crates/forge/src/lib.rs index 9e5067ed89..427ef1a59e 100644 --- a/crates/forge/src/lib.rs +++ b/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, }, /// Create a new Forge project at 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. + /// 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, + /// Additional arguments for cairo-coverage or cairo-profiler #[arg(last = true)] additional_args: Vec, @@ -327,7 +334,7 @@ pub fn main_execution(ui: Arc) -> Result { .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)?; diff --git a/crates/forge/src/partition.rs b/crates/forge/src/partition.rs new file mode 100644 index 0000000000..e32e918bd7 --- /dev/null +++ b/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 { + let parts: Vec<&str> = s.split('/').collect(); + + if parts.len() != 2 { + return Err("Partition must be in the format /".to_string()); + } + + let index = parts[0] + .parse::() + .map_err(|_| "INDEX must be a positive integer".to_string())?; + let total = parts[1] + .parse::() + .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); + +#[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::().unwrap(); + assert_eq!(partition.index, 2); + assert_eq!(partition.total, 5); + } + + #[test] + fn test_invalid_format() { + let err = "2-5".parse::().unwrap_err(); + assert_eq!(err, "Partition must be in the format /"); + } + + // #[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::().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::().unwrap_err(); + assert_eq!(err, "Invalid partition values: ensure 1 <= INDEX <= TOTAL"); + } +} diff --git a/crates/forge/tests/e2e/mod.rs b/crates/forge/tests/e2e/mod.rs index cfa90120fe..8a278f696f 100644 --- a/crates/forge/tests/e2e/mod.rs +++ b/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; diff --git a/crates/forge/tests/e2e/partitioning.rs b/crates/forge/tests/e2e/partitioning.rs new file mode 100644 index 0000000000..c82c938066 --- /dev/null +++ b/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 ' cannot be used with '--exact' + "}, + ); +}