add test that succeeds on retry (#201)

## Type

- [ ] feature
- [ ] bugfix
- [ ] dev (no functional changes, no API changes)
- [ ] fmt (formatting, renaming)
- [ ] build
- [ ] docs
- [ ] testing

## Description


## Breaking changes?

- [ ] yes
- [ ] no

<!-- Reviewable:start -->
- - -
This change is [<img src="https://reviewable.io/review_button.svg" height="34" align="absmiddle" alt="Reviewable"/>](https://reviewable.io/reviews/starkware-libs/bootloader-hints/201)
<!-- Reviewable:end -->

Author: nitsan-starkware

Cargo.toml

@@ -42,3 +42,4 @@ tracing = "0.1.40"
 sonic-rs = "0.3.17"
 starknet-ff = "0.3.7"
 mockall = "0.13.1"
+rstest = "0.21"

crates/stwo_run_and_prove/Cargo.toml

@@ -21,3 +21,4 @@ tracing.workspace = true
 [dev-dependencies]
 tempfile.workspace = true
 mockall.workspace = true
+rstest.workspace = true

crates/stwo_run_and_prove/src/main.rs

@@ -239,7 +239,7 @@ fn prove_with_retries(
     std::fs::create_dir_all(&prove_config.proofs_dir)?;
     let proof_format = prove_config.proof_format;
 
-    for i in 1..prove_config.n_proof_attempts + 1 {
+    for i in 1..=prove_config.n_proof_attempts + 1 {
         info!(
             "Attempting to generate proof {}/{}.",
             i, prove_config.n_proof_attempts
@@ -414,6 +414,7 @@ fn save_output_to_file(
 #[cfg(test)]
 mod tests {
     use super::*;
+    use rstest::rstest;
     use std::fs;
     use tempfile::{NamedTempFile, TempDir, TempPath};
 
@@ -522,6 +523,37 @@ mod tests {
         (program_output_tempfile, proofs_tempdir)
     }
 
+    fn run_with_mock_prover_succeeds_on_retry(n_proof_attempts: usize) -> (TempPath, TempDir) {
+        let (args, program_output_tempfile, proofs_tempdir) = prepare_args(n_proof_attempts);
+        let mut mock_prover = Box::new(MockProverTrait::new());
+
+        // Create iterator that return errors for the first n-1 attempts, and a successful result
+        // for the last attempt.
+        let mut results = (0..n_proof_attempts.saturating_sub(1))
+            .map(|_| Err(StwoRunAndProveError::Verification))
+            .chain(std::iter::once(Ok(vec![ARRAY_SUM_EXPECTED_OUTPUT])));
+
+        mock_prover
+            .expect_choose_channel_and_prove()
+            .times(n_proof_attempts)
+            .returning(move |_, proof_file, _, _, _| {
+                let expected_proof_file = get_path(EXPECTED_PROOF_FILE_NAME);
+                fs::copy(&expected_proof_file, &proof_file).expect("Failed to copy proof file.");
+                results.next().unwrap()
+            });
+
+        let successful_proof_attempt =
+            run_stwo_run_and_prove(args, mock_prover).expect("failed to run stwo_run_and_prove");
+
+        assert_eq!(
+            successful_proof_attempt, n_proof_attempts,
+            "successful proof attempt should be {:?}, but got {:?}",
+            n_proof_attempts, successful_proof_attempt
+        );
+
+        (program_output_tempfile, proofs_tempdir)
+    }
+
     #[test]
     fn test_stwo_run_and_prove() {
         let (output_temp_file, proofs_temp_dir) = run_with_successful_mock_prover(1);
@@ -553,25 +585,50 @@ mod tests {
     }
 
     #[test]
-    fn test_stwo_run_and_prove_retries() {
+    fn test_stwo_run_and_prove_all_retries_fail() {
         let n_proof_attempts = 3;
         let (output_temp_file, proofs_temp_dir) =
             run_with_verification_error_mock_prover(n_proof_attempts);
         let proofs_dir = proofs_temp_dir.path().to_path_buf();
 
-        for i in 1..n_proof_attempts + 1 {
+        (1..=n_proof_attempts).for_each(|i| {
             let proof_file = proofs_dir.join(format!("proof_{}", i));
             assert!(
                 proof_file.exists(),
                 "Proof file {:?} should exist after running with verifier failures",
                 i,
             );
-        }
+        });
+
         assert!(
             is_file_empty(&output_temp_file.to_path_buf()).unwrap(),
             "Output file should be empty after running with verifier failures",
         );
     }
+
+    #[rstest]
+    #[case::two_tries(2)]
+    #[case::three_tries(3)]
+    fn test_stwo_run_and_prove_succeeds_on_retry(#[case] n_proof_attempts: usize) {
+        let (output_temp_file, proofs_temp_dir) =
+            run_with_mock_prover_succeeds_on_retry(n_proof_attempts);
+        let proofs_dir = proofs_temp_dir.path().to_path_buf();
+
+        (1..=n_proof_attempts).for_each(|i| {
+            let proof_file = proofs_dir.join(format!("proof_{}", i));
+            assert!(
+                proof_file.exists(),
+                "Proof file {:?} should exist after a run that succeeds on attempt {:?} of proof and verify",
+                i, n_proof_attempts,
+            );
+        });
+
+        assert!(
+            !is_file_empty(&output_temp_file.to_path_buf()).unwrap(),
+            "Output file should not be empty after a run that succeeds on attempt {:?} of proof and verify",
+            n_proof_attempts,
+        );
+    }
 }
 
 // TODO(nitsan): Tests -