From 8fe7a9ff8cf12ef5a97e8de87a7952b5264b5607 Mon Sep 17 00:00:00 2001
From: Tamir Hemo <tamir@succinct.xyz>
Date: Tue, 20 Feb 2024 20:49:10 -0800
Subject: [PATCH] fix: verify proof in tests (#280)

---
 core/src/syscall/precompiles/keccak256/mod.rs | 19 ++++++-------------
 .../syscall/precompiles/sha256/extend/mod.rs  | 18 ++++--------------
 2 files changed, 10 insertions(+), 27 deletions(-)

diff --git a/core/src/syscall/precompiles/keccak256/mod.rs b/core/src/syscall/precompiles/keccak256/mod.rs
index 3c3b018eb9..3163b4ff1e 100644
--- a/core/src/syscall/precompiles/keccak256/mod.rs
+++ b/core/src/syscall/precompiles/keccak256/mod.rs
@@ -45,11 +45,10 @@ impl KeccakPermuteChip {
 
 #[cfg(test)]
 pub mod permute_tests {
+    use crate::utils::run_test;
     use crate::{
         runtime::{Instruction, Opcode, Program, Runtime},
-        stark::{LocalProver, RiscvStark},
-        utils::{self, tests::KECCAK_PERMUTE_ELF, BabyBearPoseidon2, StarkUtils},
-        SP1Prover, SP1Stdin,
+        utils::{self, tests::KECCAK_PERMUTE_ELF},
     };
 
     pub fn keccak_permute_program() -> Program {
@@ -78,23 +77,17 @@ pub mod permute_tests {
     }
 
     #[test]
-    fn prove_babybear() {
+    fn test_keccak_permute_prove_babybear() {
         utils::setup_logger();
-        let config = BabyBearPoseidon2::new();
-        let mut challenger = config.challenger();
 
         let program = keccak_permute_program();
-        let mut runtime = Runtime::new(program);
-        runtime.run();
-
-        let machine = RiscvStark::new(config);
-        let (pk, _) = machine.setup(runtime.program.as_ref());
-        machine.prove::<LocalProver<_>>(&pk, &mut runtime.record, &mut challenger);
+        run_test(program).unwrap();
     }
 
     #[test]
     fn test_keccak_permute_program_prove() {
         utils::setup_logger();
-        SP1Prover::prove(KECCAK_PERMUTE_ELF, SP1Stdin::new()).unwrap();
+        let program = Program::from(KECCAK_PERMUTE_ELF);
+        run_test(program).unwrap();
     }
 }
diff --git a/core/src/syscall/precompiles/sha256/extend/mod.rs b/core/src/syscall/precompiles/sha256/extend/mod.rs
index 1671824067..8238ed647b 100644
--- a/core/src/syscall/precompiles/sha256/extend/mod.rs
+++ b/core/src/syscall/precompiles/sha256/extend/mod.rs
@@ -47,9 +47,8 @@ pub mod extend_tests {
     use crate::{
         air::MachineAir,
         alu::AluEvent,
-        runtime::{ExecutionRecord, Instruction, Opcode, Program, Runtime},
-        stark::{LocalProver, RiscvStark},
-        utils::{BabyBearPoseidon2, StarkUtils},
+        runtime::{ExecutionRecord, Instruction, Opcode, Program},
+        utils::run_test,
     };
 
     use super::ShaExtendChip;
@@ -82,17 +81,8 @@ pub mod extend_tests {
     }
 
     #[test]
-    fn prove_babybear() {
-        let config = BabyBearPoseidon2::new();
-        let mut challenger = config.challenger();
-
+    fn test_sha_prove() {
         let program = sha_extend_program();
-        let mut runtime = Runtime::new(program);
-        runtime.run();
-
-        let machine = RiscvStark::new(config);
-
-        let (pk, _) = machine.setup(runtime.program.as_ref());
-        machine.prove::<LocalProver<_>>(&pk, &mut runtime.record, &mut challenger);
+        run_test(program).unwrap();
     }
 }