test: VirtIO PCI device creation and restoration

Do a bit of refactoring on the test code that inserts VirtIO devices in
a Vmm object and then add a test for PCI devices which creates a Vmm
with PCI devices and then restores them and makes sure that everything
is restored as expected.

Signed-off-by: Babis Chalios <[email protected]>

Author: bchalios

src/vmm/src/builder.rs

@@ -848,7 +848,6 @@ pub(crate) mod tests {
 
         assert!(
             vmm.device_manager
-                .mmio_devices
                 .get_virtio_device(TYPE_VSOCK, &vsock_dev_id)
                 .is_some()
         );
@@ -874,7 +873,6 @@ pub(crate) mod tests {
 
         assert!(
             vmm.device_manager
-                .mmio_devices
                 .get_virtio_device(TYPE_RNG, ENTROPY_DEV_ID)
                 .is_some()
         );
@@ -909,7 +907,6 @@ pub(crate) mod tests {
 
         assert!(
             vmm.device_manager
-                .mmio_devices
                 .get_virtio_device(TYPE_BALLOON, BALLOON_DEV_ID)
                 .is_some()
         );
@@ -961,7 +958,6 @@ pub(crate) mod tests {
             assert!(cmdline_contains(&cmdline, "root=/dev/vda ro"));
             assert!(
                 vmm.device_manager
-                    .mmio_devices
                     .get_virtio_device(TYPE_BLOCK, drive_id.as_str())
                     .is_some()
             );
@@ -983,7 +979,6 @@ pub(crate) mod tests {
             assert!(cmdline_contains(&cmdline, "root=PARTUUID=0eaa91a0-01 rw"));
             assert!(
                 vmm.device_manager
-                    .mmio_devices
                     .get_virtio_device(TYPE_BLOCK, drive_id.as_str())
                     .is_some()
             );
@@ -1006,7 +1001,6 @@ pub(crate) mod tests {
             assert!(!cmdline_contains(&cmdline, "root=/dev/vda"));
             assert!(
                 vmm.device_manager
-                    .mmio_devices
                     .get_virtio_device(TYPE_BLOCK, drive_id.as_str())
                     .is_some()
             );
@@ -1044,19 +1038,16 @@ pub(crate) mod tests {
             assert!(cmdline_contains(&cmdline, "root=PARTUUID=0eaa91a0-01 rw"));
             assert!(
                 vmm.device_manager
-                    .mmio_devices
                     .get_virtio_device(TYPE_BLOCK, "root")
                     .is_some()
             );
             assert!(
                 vmm.device_manager
-                    .mmio_devices
                     .get_virtio_device(TYPE_BLOCK, "secondary")
                     .is_some()
             );
             assert!(
                 vmm.device_manager
-                    .mmio_devices
                     .get_virtio_device(TYPE_BLOCK, "third")
                     .is_some()
             );
@@ -1086,7 +1077,6 @@ pub(crate) mod tests {
             assert!(cmdline_contains(&cmdline, "root=/dev/vda rw"));
             assert!(
                 vmm.device_manager
-                    .mmio_devices
                     .get_virtio_device(TYPE_BLOCK, drive_id.as_str())
                     .is_some()
             );
@@ -1108,7 +1098,6 @@ pub(crate) mod tests {
             assert!(cmdline_contains(&cmdline, "root=PARTUUID=0eaa91a0-01 ro"));
             assert!(
                 vmm.device_manager
-                    .mmio_devices
                     .get_virtio_device(TYPE_BLOCK, drive_id.as_str())
                     .is_some()
             );
@@ -1130,7 +1119,6 @@ pub(crate) mod tests {
             assert!(cmdline_contains(&cmdline, "root=/dev/vda rw"));
             assert!(
                 vmm.device_manager
-                    .mmio_devices
                     .get_virtio_device(TYPE_BLOCK, drive_id.as_str())
                     .is_some()
             );

src/vmm/src/device_manager/mod.rs

@@ -375,6 +375,36 @@ impl DeviceManager {
             Self::do_mark_virtio_queue_memory_dirty(virtio_device, mem);
         }
     }
+
+    /// Get a VirtIO device of type `virtio_type` with ID `device_id`
+    pub fn get_virtio_device(
+        &self,
+        virtio_type: u32,
+        device_id: &str,
+    ) -> Option<Arc<Mutex<dyn VirtioDevice>>> {
+        if self.pci_devices.pci_segment.is_some() {
+            let pci_device = self.pci_devices.get_virtio_device(virtio_type, device_id)?;
+            Some(
+                pci_device
+                    .lock()
+                    .expect("Poisoned lock")
+                    .virtio_device()
+                    .clone(),
+            )
+        } else {
+            let mmio_device = self
+                .mmio_devices
+                .get_virtio_device(virtio_type, device_id)?;
+            Some(
+                mmio_device
+                    .inner
+                    .lock()
+                    .expect("Poisoned lock")
+                    .device()
+                    .clone(),
+            )
+        }
+    }
 }
 
 #[derive(Debug, Default, Clone, Serialize, Deserialize)]