From a24d877f969698da5afaa57119a061808b1355b4 Mon Sep 17 00:00:00 2001
From: Purna Pavan Chandra <paekkaladevi@linux.microsoft.com>
Date: Tue, 23 Apr 2024 06:41:02 +0000
Subject: [PATCH] vmm: Add unit tests for restore

The tests provide different inputs to parse and validate functions to
make sure parsing and error handling is as per expectation.

Signed-off-by: Purna Pavan Chandra <paekkaladevi@linux.microsoft.com>
---
 vmm/src/config.rs | 138 ++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 138 insertions(+)

diff --git a/vmm/src/config.rs b/vmm/src/config.rs
index b90feeddfc..b1887c02ec 100644
--- a/vmm/src/config.rs
+++ b/vmm/src/config.rs
@@ -3535,6 +3535,144 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn test_restore_parsing() -> Result<()> {
+        assert_eq!(
+            RestoreConfig::parse("source_url=/path/to/snapshot")?,
+            RestoreConfig {
+                source_url: PathBuf::from("/path/to/snapshot"),
+                prefault: false,
+                net_ids: None,
+                net_fds: None,
+            }
+        );
+        assert_eq!(
+            RestoreConfig::parse(
+                "source_url=/path/to/snapshot,prefault=off,net_ids=[net0,net1],net_fds=[3,4,5,6]"
+            )?,
+            RestoreConfig {
+                source_url: PathBuf::from("/path/to/snapshot"),
+                prefault: false,
+                net_ids: Some(vec!["net0".to_string(), "net1".to_string()]),
+                net_fds: Some(vec![3, 4, 5, 6]),
+            }
+        );
+        // Parsing should fail as source_url is a required field
+        assert!(!RestoreConfig::parse("prefault=off").is_ok());
+        Ok(())
+    }
+
+    #[test]
+    fn test_restore_config_validation() {
+        // interested in only VmConfig.net, so set rest to default values
+        let mut snapshot_vm_config = VmConfig {
+            cpus: CpusConfig::default(),
+            memory: MemoryConfig::default(),
+            payload: None,
+            rate_limit_groups: None,
+            disks: None,
+            rng: RngConfig::default(),
+            balloon: None,
+            fs: None,
+            pmem: None,
+            serial: default_serial(),
+            console: default_console(),
+            debug_console: DebugConsoleConfig::default(),
+            devices: None,
+            user_devices: None,
+            vdpa: None,
+            vsock: None,
+            pvpanic: false,
+            iommu: false,
+            sgx_epc: None,
+            numa: None,
+            watchdog: false,
+            platform: None,
+            tpm: None,
+            preserved_fds: None,
+            net: Some(vec![
+                NetConfig {
+                    id: Some("net0".to_owned()),
+                    num_queues: 2,
+                    fds: Some(vec![-1, -1, -1, -1]),
+                    ..net_fixture()
+                },
+                NetConfig {
+                    id: Some("net1".to_owned()),
+                    num_queues: 1,
+                    fds: Some(vec![-1, -1]),
+                    ..net_fixture()
+                },
+                NetConfig {
+                    id: Some("net2".to_owned()),
+                    fds: None,
+                    ..net_fixture()
+                },
+            ]),
+        };
+
+        let valid_config = RestoreConfig {
+            source_url: PathBuf::from("/path/to/snapshot"),
+            prefault: false,
+            net_ids: Some(vec!["net0".to_string(), "net1".to_string()]),
+            net_fds: Some(vec![3, 4, 5, 6, 7, 8]),
+        };
+        assert!(valid_config.validate(&snapshot_vm_config).is_ok());
+
+        let mut invalid_config = valid_config.clone();
+        invalid_config.net_ids = Some(vec!["netx".to_string()]);
+        assert_eq!(
+            invalid_config.validate(&snapshot_vm_config),
+            Err(ValidationError::InvalidIdentifier("netx".to_string()))
+        );
+
+        invalid_config.net_ids = Some(vec!["net0".to_string(), "net0".to_string()]);
+        assert_eq!(
+            invalid_config.validate(&snapshot_vm_config),
+            Err(ValidationError::IdentifierNotUnique("net0".to_string()))
+        );
+
+        invalid_config.net_ids = Some(vec!["net1".to_string()]);
+        assert_eq!(
+            invalid_config.validate(&snapshot_vm_config),
+            Err(ValidationError::RestoreMissingRequiredNetId(
+                "net0".to_string()
+            ))
+        );
+
+        invalid_config.net_ids = Some(vec![
+            "net0".to_string(),
+            "net1".to_string(),
+            "net2".to_string(),
+        ]);
+        assert_eq!(
+            invalid_config.validate(&snapshot_vm_config),
+            Err(ValidationError::RestoreNonFdNetIdNotExpected(
+                "net2".to_string()
+            ))
+        );
+
+        invalid_config.net_ids = Some(vec!["net0".to_string(), "net1".to_string()]);
+        invalid_config.net_fds = Some(vec![3, 4, 5, 6]);
+        assert_eq!(
+            invalid_config.validate(&snapshot_vm_config),
+            Err(ValidationError::RestoreNetFdCountMismatch(4, 6))
+        );
+
+        let another_valid_config = RestoreConfig {
+            source_url: PathBuf::from("/path/to/snapshot"),
+            prefault: false,
+            net_ids: None,
+            net_fds: None,
+        };
+        snapshot_vm_config.net = Some(vec![NetConfig {
+            id: Some("net2".to_owned()),
+            fds: None,
+            ..net_fixture()
+        }]);
+        assert!(another_valid_config.validate(&snapshot_vm_config).is_ok());
+    }
+
     fn platform_fixture() -> PlatformConfig {
         PlatformConfig {
             num_pci_segments: MAX_NUM_PCI_SEGMENTS,