feat: Handle should_fail_with case (#2541)

Co-authored-by: sirasistant <sirasistant@gmail.com>
Co-authored-by: Maxim Vezenov <mvezenov@gmail.com>

crates/nargo/src/errors.rs

@@ -17,6 +17,32 @@ pub enum NargoError {
     ForeignCallError(#[from] ForeignCallError),
 }
 
+impl NargoError {
+    /// Extracts the user defined failure message from the ExecutionError
+    /// If one exists.
+    ///
+    /// We want to extract the user defined error so that we can compare it
+    /// in tests to expected failure messages
+    pub fn user_defined_failure_message(&self) -> Option<&str> {
+        let execution_error = match self {
+            NargoError::ExecutionError(error) => error,
+            _ => return None,
+        };
+
+        match execution_error {
+            ExecutionError::AssertionFailed(message, _) => Some(message),
+            ExecutionError::SolvingError(error) => match error {
+                OpcodeResolutionError::IndexOutOfBounds { .. }
+                | OpcodeResolutionError::UnsupportedBlackBoxFunc(_)
+                | OpcodeResolutionError::OpcodeNotSolvable(_)
+                | OpcodeResolutionError::UnsatisfiedConstrain { .. } => None,
+                OpcodeResolutionError::BrilligFunctionFailed { message, .. } => Some(message),
+                OpcodeResolutionError::BlackBoxFunctionFailed(_, reason) => Some(reason),
+            },
+        }
+    }
+}
+
 #[derive(Debug, Error)]
 pub enum ExecutionError {
     #[error("Failed assertion: '{}'", .0)]

crates/nargo/src/ops/test.rs

@@ -3,6 +3,8 @@ use noirc_driver::{compile_no_check, CompileOptions};
 use noirc_errors::FileDiagnostic;
 use noirc_frontend::hir::{def_map::TestFunction, Context};
 
+use crate::NargoError;
+
 use super::execute_circuit;
 
 pub enum TestStatus {
@@ -25,44 +27,97 @@ pub fn run_test<B: BlackBoxFunctionSolver>(
             // otherwise constraints involving these expressions will not error.
             let circuit_execution =
                 execute_circuit(blackbox_solver, program.circuit, WitnessMap::new(), show_output);
+            test_status_program_compile_pass(test_function, circuit_execution)
+        }
+        Err(diag) => test_status_program_compile_fail(diag, test_function),
+    }
+}
+
+/// Test function failed to compile
+///
+/// Note: This could be because the compiler was able to deduce
+/// that a constraint was never satisfiable.
+/// An example of this is the program `assert(false)`
+/// In that case, we check if the test function should fail, and if so, we return `TestStatus::Pass`.
+fn test_status_program_compile_fail(
+    diag: FileDiagnostic,
+    test_function: TestFunction,
+) -> TestStatus {
+    // The test has failed compilation, but it should never fail. Report error.
+    if !test_function.should_fail() {
+        return TestStatus::CompileError(diag);
+    }
+
+    // The test has failed compilation, check if it is because the program is never satisfiable.
+    // If it is never satisfiable, then this is the expected behavior.
+    let program_is_never_satisfiable = diag.diagnostic.message.contains("Failed constraint");
+    if !program_is_never_satisfiable {
+        // The test has failed compilation, but its a compilation error. Report error
+        return TestStatus::CompileError(diag);
+    }
+
+    check_expected_failure_message(test_function, &diag.diagnostic.message)
+}
 
+/// The test function compiled successfully.
+///
+/// We now check whether execution passed/failed and whether it should have
+/// passed/failed to determine the test status.
+fn test_status_program_compile_pass(
+    test_function: TestFunction,
+    circuit_execution: Result<WitnessMap, NargoError>,
+) -> TestStatus {
+    let circuit_execution_err = match circuit_execution {
+        // Circuit execution was successful; ie no errors or unsatisfied constraints
+        // were encountered.
+        Ok(_) => {
             if test_function.should_fail() {
-                match circuit_execution {
-                    Ok(_) => TestStatus::Fail {
-                        // TODO: Improve color variations on this message
-                        message: "error: Test passed when it should have failed".to_string(),
-                    },
-                    Err(_) => TestStatus::Pass,
-                }
-            } else {
-                match circuit_execution {
-                    Ok(_) => TestStatus::Pass,
-                    Err(error) => TestStatus::Fail { message: error.to_string() },
-                }
+                return TestStatus::Fail {
+                    message: "error: Test passed when it should have failed".to_string(),
+                };
             }
+            return TestStatus::Pass;
         }
-        // Test function failed to compile
-        //
-        // Note: This could be because the compiler was able to deduce
-        // that a constraint was never satisfiable.
-        // An example of this is the program `assert(false)`
-        // In that case, we check if the test function should fail, and if so, we return `TestStatus::Pass`.
-        Err(diag) => {
-            // The test has failed compilation, but it should never fail. Report error.
-            if !test_function.should_fail() {
-                return TestStatus::CompileError(diag);
-            }
+        Err(err) => err,
+    };
 
-            // The test has failed compilation, check if it is because the program is never satisfiable.
-            // If it is never satisfiable, then this is the expected behavior.
-            let program_is_never_satisfiable =
-                diag.diagnostic.message.contains("Failed constraint");
-            if program_is_never_satisfiable {
-                return TestStatus::Pass;
-            }
+    // If we reach here, then the circuit execution failed.
+    //
+    // Check if the function should have passed
+    let test_should_have_passed = !test_function.should_fail();
+    if test_should_have_passed {
+        return TestStatus::Fail { message: circuit_execution_err.to_string() };
+    }
 
-            // The test has failed compilation, but its a compilation error. Report error
-            TestStatus::CompileError(diag)
-        }
+    check_expected_failure_message(
+        test_function,
+        circuit_execution_err.user_defined_failure_message().unwrap_or_default(),
+    )
+}
+
+fn check_expected_failure_message(test_function: TestFunction, got_error: &str) -> TestStatus {
+    // Extract the expected failure message, if there was one
+    //
+    // #[test(should_fail)] will not produce any message
+    // #[test(should_fail_with = "reason")] will produce a message
+    //
+    let expected_failure_message = match test_function.failure_reason() {
+        Some(reason) => reason,
+        None => return TestStatus::Pass,
+    };
+
+    let expected_failure_message_matches =
+        got_error == expected_failure_message;
+    if expected_failure_message_matches {
+        return TestStatus::Pass;
     }
+
+    // The expected failure message does not match the actual failure message
+    return TestStatus::Fail {
+        message: format!(
+            "\nerror: Test failed with the wrong message. \nExpected: {} \nGot: {}",
+            test_function.failure_reason().unwrap_or_default(),
+            got_error.trim_matches('\'')
+        ),
+    };
 }

crates/noirc_frontend/src/hir/def_map/mod.rs

@@ -251,10 +251,20 @@ impl TestFunction {
 
     /// Returns true if the test function has been specified to fail
     /// This is done by annotating the function with `#[test(should_fail)]`
+    /// or `#[test(should_fail_with = "reason")]`
     pub fn should_fail(&self) -> bool {
         match self.scope {
-            TestScope::ShouldFail => true,
+            TestScope::ShouldFailWith { .. } => true,
             TestScope::None => false,
         }
     }
+
+    /// Returns the reason for the test function to fail if specified
+    /// by the user.
+    pub fn failure_reason(&self) -> Option<&str> {
+        match &self.scope {
+            TestScope::None => None,
+            TestScope::ShouldFailWith { reason } => reason.as_deref(),
+        }
+    }
 }

crates/noirc_frontend/src/lexer/lexer.rs

@@ -532,7 +532,24 @@ mod tests {
         let mut lexer = Lexer::new(input);
 
         let token = lexer.next().unwrap().unwrap();
-        assert_eq!(token.token(), &Token::Attribute(Attribute::Test(TestScope::ShouldFail)));
+        assert_eq!(
+            token.token(),
+            &Token::Attribute(Attribute::Test(TestScope::ShouldFailWith { reason: None }))
+        );
+    }
+
+    #[test]
+    fn test_attribute_with_valid_scope_should_fail_with() {
+        let input = r#"#[test(should_fail_with = "hello")]"#;
+        let mut lexer = Lexer::new(input);
+
+        let token = lexer.next().unwrap().unwrap();
+        assert_eq!(
+            token.token(),
+            &Token::Attribute(Attribute::Test(TestScope::ShouldFailWith {
+                reason: Some("hello".to_owned())
+            }))
+        );
     }
 
     #[test]

crates/noirc_frontend/src/lexer/token.rs

@@ -319,30 +319,42 @@ impl IntType {
 /// TestScope is used to specify additional annotations for test functions
 #[derive(PartialEq, Eq, Hash, Debug, Clone, PartialOrd, Ord)]
 pub enum TestScope {
-    /// If a test has a scope of ShouldFail, then it is expected to fail
-    ShouldFail,
+    /// If a test has a scope of ShouldFailWith, then it can only pass
+    /// if it fails with the specified reason. If the reason is None, then
+    /// the test must unconditionally fail
+    ShouldFailWith { reason: Option<String> },
     /// No scope is applied and so the test must pass
     None,
 }
 
 impl TestScope {
     fn lookup_str(string: &str) -> Option<TestScope> {
-        match string {
-            "should_fail" => Some(TestScope::ShouldFail),
+        match string.trim() {
+            "should_fail" => Some(TestScope::ShouldFailWith { reason: None }),
+            s if s.starts_with("should_fail_with") => {
+                let parts: Vec<&str> = s.splitn(2, '=').collect();
+                if parts.len() == 2 {
+                    let reason = parts[1].trim();
+                    let reason = reason.trim_matches('"');
+                    Some(TestScope::ShouldFailWith { reason: Some(reason.to_string()) })
+                } else {
+                    None
+                }
+            }
             _ => None,
         }
     }
-    fn as_str(&self) -> &'static str {
-        match self {
-            TestScope::ShouldFail => "should_fail",
-            TestScope::None => "",
-        }
-    }
 }
 
 impl fmt::Display for TestScope {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        write!(f, "({})", self.as_str())
+        match self {
+            TestScope::None => write!(f, ""),
+            TestScope::ShouldFailWith { reason } => match reason {
+                Some(failure_reason) => write!(f, "(should_fail_with = ({}))", failure_reason),
+                None => write!(f, "should_fail"),
+            },
+        }
     }
 }
 
@@ -391,6 +403,9 @@ impl Attribute {
                         || ch == '_'
                         || ch == '('
                         || ch == ')'
+                        || ch == '='
+                        || ch == '"'
+                        || ch == ' '
                 })
                 .then_some(());