fix: adjust test

Signed-off-by: Tomasz Andrzejak <andreiltd@gmail.com>

Author: andreiltd

src/hyperlight_common/src/virtq/consumer.rs

@@ -1406,6 +1406,11 @@ mod tests {
         // The rejected chain is returned. The earlier chain remains in flight.
         assert_eq!(producer.poll().unwrap().unwrap().token(), rejected);
         assert!(producer.poll().unwrap().is_none());
+
+        drop(consumer);
+        // SAFETY: The peer is stopped and all consumer chain handles are dropped.
+        unsafe { producer.reset().unwrap() };
+        assert_eq!(producer.pool().num_live(), 0);
     }
 
     #[test]
@@ -1428,6 +1433,11 @@ mod tests {
         assert_eq!(consumer.inner.num_inflight(), 0);
         assert_eq!(consumer.next_token, 0);
         assert!(producer.poll().unwrap().is_none());
+
+        drop(consumer);
+        // SAFETY: The peer is stopped and all consumer chain handles are dropped.
+        unsafe { producer.reset().unwrap() };
+        assert_eq!(producer.pool().num_live(), 0);
     }
 
     #[test]
@@ -1452,6 +1462,11 @@ mod tests {
         assert_eq!(consumer.inner.num_inflight(), 1);
         assert_eq!(consumer.next_token, 1);
         assert!(producer.poll().unwrap().is_none());
+
+        drop(consumer);
+        // SAFETY: The peer is stopped and all consumer chain handles are dropped.
+        unsafe { producer.reset().unwrap() };
+        assert_eq!(producer.pool().num_live(), 0);
     }
 
     #[test]

src/hyperlight_host/src/mem/mgr.rs

@@ -1027,12 +1027,14 @@ mod tests {
     use crate::sandbox::SandboxConfiguration;
 
     fn manager(case: &TestCase) -> SandboxMemoryManager<HostSharedMemory> {
+        let host_page_size = page_size::get();
+
         #[cfg(not(unshared_snapshot_mem))]
         let shared_mem =
-            ReadonlySharedMemory::from_bytes(&vec![0; vmem::PAGE_SIZE], vmem::PAGE_SIZE).unwrap();
+            ReadonlySharedMemory::from_bytes(&vec![0; host_page_size], host_page_size).unwrap();
 
         #[cfg(unshared_snapshot_mem)]
-        let shared_mem = ExclusiveSharedMemory::new(vmem::PAGE_SIZE)
+        let shared_mem = ExclusiveSharedMemory::new(host_page_size)
             .unwrap()
             .build()
             .0;

src/hyperlight_host/src/sandbox/initialized_multi_use.rs

@@ -4920,6 +4920,9 @@ mod tests {
             let mut sbox = make_sandbox();
             sbox.call::<i32>("AddToStatic", 1i32).unwrap();
 
+            // Checkpointing repairs the rings, so inject corruption afterward.
+            sbox.checkpoint_transport_for_snapshot().unwrap();
+
             let layout = &sbox.mem_mgr.layout;
             let scratch_base = layout::scratch_base_gpa(layout.get_scratch_size());