diff --git a/src/lib.rs b/src/lib.rs
index fd84d93..35bc576 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -33,7 +33,7 @@ mod tests {
     #[test]
     /// Tests whether two numbers in memory can be added together
     /// in the ZKVM
-    fn test_preflight_1() {
+    fn test_preflight_add_memory() {
         let instructions = vec![
             Instruction::Lb(Register::R0, MemoryLocation(0x40)),
             Instruction::Lb(Register::R1, MemoryLocation(0x41)),
@@ -63,6 +63,14 @@ mod tests {
         assert!(simulation.is_ok());
         let simulation = simulation.unwrap();
 
-        println!("{:#?}", simulation);
+        assert_eq!(
+            simulation.trace_rows[simulation
+                .trace_rows
+                .len()
+                - 1]
+            .get_memory_at(&expected.0)
+            .unwrap(),
+            expected.1
+        );
     }
 }
diff --git a/src/preflight_simulator.rs b/src/preflight_simulator.rs
index 571f1a6..23ee7f4 100644
--- a/src/preflight_simulator.rs
+++ b/src/preflight_simulator.rs
@@ -90,8 +90,7 @@ impl SimulationRow {
             .context("instruction not found")?;
 
         let is_halted = instruction == Instruction::Halt;
-        let mut registers = self
-            .registers;
+        let mut registers = self.registers;
 
         let mut memory_snapshot = self
             .memory_snapshot
@@ -148,6 +147,20 @@ impl SimulationRow {
             memory_snapshot,
         })
     }
+
+    pub fn get_memory_at(
+        &self,
+        address: &u8,
+    ) -> Option<u8> {
+        self.memory_snapshot
+            .get(address)
+            .copied()
+    }
+
+    pub fn get_registers(&self) -> [u8; REGISTER_COUNT] {
+        self.registers
+            .clone()
+    }
 }
 
 /// Unconstrainted Preflight Simulation of the program built