diff --git a/ceno_emul/src/tracer.rs b/ceno_emul/src/tracer.rs
index 807089618..6f9d623f7 100644
--- a/ceno_emul/src/tracer.rs
+++ b/ceno_emul/src/tracer.rs
@@ -116,12 +116,12 @@ impl StepRecord {
     pub fn new_u_instruction(
         cycle: Cycle,
         pc: ByteAddr,
-        insn_code: Word,
+        insn_code: u32,
         rd: Change<Word>,
         prev_cycle: Cycle,
     ) -> StepRecord {
         let pc = Change::new(pc, pc + PC_STEP_SIZE);
-        StepRecord::new_insn(cycle, pc, insn_code, None, None, Some(rd), prev_cycle)
+        StepRecord::new_insn2(cycle, pc, insn_code, None, None, Some(rd), prev_cycle)
     }
 
     pub fn new_j_instruction(
diff --git a/ceno_zkvm/src/instructions/riscv/jump/test.rs b/ceno_zkvm/src/instructions/riscv/jump/test.rs
index 9ac13cf44..8db48d7b4 100644
--- a/ceno_zkvm/src/instructions/riscv/jump/test.rs
+++ b/ceno_zkvm/src/instructions/riscv/jump/test.rs
@@ -1,4 +1,4 @@
-use ceno_emul::{ByteAddr, Change, PC_STEP_SIZE, StepRecord};
+use ceno_emul::{ByteAddr, Change, EncodedInstruction, InsnKind, PC_STEP_SIZE, StepRecord};
 use goldilocks::GoldilocksExt2;
 use itertools::Itertools;
 use multilinear_extensions::mle::IntoMLEs;
@@ -6,7 +6,7 @@ use multilinear_extensions::mle::IntoMLEs;
 use crate::{
     circuit_builder::{CircuitBuilder, ConstraintSystem},
     instructions::Instruction,
-    scheme::mock_prover::{MOCK_PC_AUIPC, MOCK_PC_JAL, MOCK_PC_LUI, MOCK_PROGRAM, MockProver},
+    scheme::mock_prover::{MOCK_PC_JAL, MOCK_PC_START, MOCK_PROGRAM, MockProver},
 };
 
 use super::{AuipcInstruction, JalInstruction, LuiInstruction};
@@ -54,6 +54,10 @@ fn test_opcode_jal() {
     );
 }
 
+fn imm(imm: u32) -> u32 {
+    // valid imm is imm[12:31] in U-type
+    imm << 12
+}
 #[test]
 fn test_opcode_lui() {
     let mut cs = ConstraintSystem::<GoldilocksExt2>::new(|| "riscv");
@@ -69,22 +73,22 @@ fn test_opcode_lui() {
         .unwrap()
         .unwrap();
 
-    let lui_insn = MOCK_PROGRAM[22];
-    let imm = lui_insn & 0xfffff000;
+    let imm_value = imm(0x90005);
+    let insn_code = EncodedInstruction::encode(InsnKind::LUI, 0, 0, 4, imm_value);
     let (raw_witin, lkm) = LuiInstruction::<GoldilocksExt2>::assign_instances(
         &config,
         cb.cs.num_witin as usize,
         vec![StepRecord::new_u_instruction(
             4,
-            MOCK_PC_LUI,
-            lui_insn,
-            Change::new(0, imm),
+            MOCK_PC_START,
+            insn_code,
+            Change::new(0, imm_value),
             0,
         )],
     )
     .unwrap();
 
-    MockProver::assert_satisfied(
+    MockProver::assert_satisfied_with_program(
         &cb,
         &raw_witin
             .de_interleaving()
@@ -92,6 +96,7 @@ fn test_opcode_lui() {
             .into_iter()
             .map(|v| v.into())
             .collect_vec(),
+        &[insn_code],
         None,
         Some(lkm),
     );
@@ -112,22 +117,22 @@ fn test_opcode_auipc() {
         .unwrap()
         .unwrap();
 
-    let auipc_insn = MOCK_PROGRAM[23];
-    let imm = auipc_insn & 0xfffff000;
+    let imm_value = imm(0x90005);
+    let insn_code = EncodedInstruction::encode(InsnKind::AUIPC, 0, 0, 4, imm_value);
     let (raw_witin, lkm) = AuipcInstruction::<GoldilocksExt2>::assign_instances(
         &config,
         cb.cs.num_witin as usize,
         vec![StepRecord::new_u_instruction(
             4,
-            MOCK_PC_AUIPC,
-            auipc_insn,
-            Change::new(0, MOCK_PC_AUIPC.0.wrapping_add(imm)),
+            MOCK_PC_START,
+            insn_code,
+            Change::new(0, MOCK_PC_START.0.wrapping_add(imm_value)),
             0,
         )],
     )
     .unwrap();
 
-    MockProver::assert_satisfied(
+    MockProver::assert_satisfied_with_program(
         &cb,
         &raw_witin
             .de_interleaving()
@@ -135,6 +140,7 @@ fn test_opcode_auipc() {
             .into_iter()
             .map(|v| v.into())
             .collect_vec(),
+        &[insn_code],
         None,
         Some(lkm),
     );