tests: Refactor test_cpu_affinity into reusable helper

Extract CPU affinity test logic from test_cpu_affinity into a
standalone _test_cpu_affinity helper that accepts a Guest
reference. The helper verifies the host has at least 4 CPUs,
boots a VM with affinity settings, and asserts vcpu0 is pinned
to cores 0,2 and vcpu1 to cores 1,3.

Add default_cpus_with_affinity_string() to Guest and
default_cpus_with_affinity() to GuestCommand in test_infra
to generate CPU arguments with affinity configuration.

Update the test_cpu_affinity call site in common_parallel to
use GuestFactory and delegate to the new helper, enabling
reuse with different guest types.

Signed-off-by: Muminul Islam <muislam@microsoft.com>

cloud-hypervisor/tests/integration.rs

@@ -2502,6 +2502,42 @@ fn _test_multi_cpu(guest: &Guest) {
     handle_child_output(r, &output);
 }
 
+fn _test_cpu_affinity(guest: &Guest) {
+    // We need the host to have at least 4 CPUs if we want to be able
+    // to run this test.
+    let host_cpus_count = exec_host_command_output("nproc");
+    assert!(
+        String::from_utf8_lossy(&host_cpus_count.stdout)
+            .trim()
+            .parse::<u16>()
+            .unwrap_or(0)
+            >= 4
+    );
+
+    let mut child = GuestCommand::new(guest)
+        .default_cpus_with_affinity()
+        .default_memory()
+        .default_kernel_cmdline()
+        .default_disks()
+        .default_net()
+        .capture_output()
+        .spawn()
+        .unwrap();
+
+    let r = std::panic::catch_unwind(|| {
+        guest.wait_vm_boot().unwrap();
+        let pid = child.id();
+        let taskset_vcpu0 = exec_host_command_output(format!("taskset -pc $(ps -T -p {pid} | grep vcpu0 | xargs | cut -f 2 -d \" \") | cut -f 6 -d \" \"").as_str());
+        assert_eq!(String::from_utf8_lossy(&taskset_vcpu0.stdout).trim(), "0,2");
+        let taskset_vcpu1 = exec_host_command_output(format!("taskset -pc $(ps -T -p {pid} | grep vcpu1 | xargs | cut -f 2 -d \" \") | cut -f 6 -d \" \"").as_str());
+        assert_eq!(String::from_utf8_lossy(&taskset_vcpu1.stdout).trim(), "1,3");
+    });
+
+    kill_child(&mut child);
+    let output = child.wait_with_output().unwrap();
+    handle_child_output(r, &output);
+}
+
 mod common_parallel {
     use std::cmp;
     use std::fs::{File, OpenOptions, copy};
@@ -2638,42 +2674,10 @@ mod common_parallel {
     #[test]
     fn test_cpu_affinity() {
         let disk_config = UbuntuDiskConfig::new(JAMMY_IMAGE_NAME.to_string());
-        let guest = Guest::new(Box::new(disk_config));
-
-        // We need the host to have at least 4 CPUs if we want to be able
-        // to run this test.
-        let host_cpus_count = exec_host_command_output("nproc");
-        assert!(
-            String::from_utf8_lossy(&host_cpus_count.stdout)
-                .trim()
-                .parse::<u16>()
-                .unwrap_or(0)
-                >= 4
-        );
-
-        let mut child = GuestCommand::new(&guest)
-            .args(["--cpus", "boot=2,affinity=[0@[0,2],1@[1,3]]"])
-            .default_memory()
-            .args(["--kernel", direct_kernel_boot_path().to_str().unwrap()])
-            .args(["--cmdline", DIRECT_KERNEL_BOOT_CMDLINE])
-            .default_disks()
-            .default_net()
-            .capture_output()
-            .spawn()
-            .unwrap();
-
-        let r = std::panic::catch_unwind(|| {
-            guest.wait_vm_boot().unwrap();
-            let pid = child.id();
-            let taskset_vcpu0 = exec_host_command_output(format!("taskset -pc $(ps -T -p {pid} | grep vcpu0 | xargs | cut -f 2 -d \" \") | cut -f 6 -d \" \"").as_str());
-            assert_eq!(String::from_utf8_lossy(&taskset_vcpu0.stdout).trim(), "0,2");
-            let taskset_vcpu1 = exec_host_command_output(format!("taskset -pc $(ps -T -p {pid} | grep vcpu1 | xargs | cut -f 2 -d \" \") | cut -f 6 -d \" \"").as_str());
-            assert_eq!(String::from_utf8_lossy(&taskset_vcpu1.stdout).trim(), "1,3");
-        });
-
-        kill_child(&mut child);
-        let output = child.wait_with_output().unwrap();
-        handle_child_output(r, &output);
+        let guest = GuestFactory::new_regular_guest_factory()
+            .create_guest(Box::new(disk_config))
+            .with_cpu(2);
+        _test_cpu_affinity(&guest);
     }
 
     #[test]

test_infra/src/lib.rs

@@ -1442,6 +1442,14 @@ impl Guest {
         )
     }
 
+    pub fn default_cpus_with_affinity_string(&self) -> String {
+        format!(
+            "boot={},affinity=[0@[0,2],1@[1,3]]{}",
+            self.num_cpu,
+            if self.nested { "" } else { ",nested=off" }
+        )
+    }
+
     pub fn default_memory_string(&self) -> String {
         format!("size={}", self.mem_size_str)
     }
@@ -1716,6 +1724,16 @@ impl<'a> GuestCommand<'a> {
         self.args(["--cpus", self.guest.default_cpus_string().as_str()])
     }
 
+    pub fn default_cpus_with_affinity(&mut self) -> &mut Self {
+        // Only support cpu affinity for 2 VCPUs for now,
+        // as it is only used in a test that validates cpu affinity is applied correctly.
+        assert_eq!(self.guest.num_cpu, 2);
+        self.args([
+            "--cpus",
+            self.guest.default_cpus_with_affinity_string().as_str(),
+        ])
+    }
+
     pub fn default_memory(&mut self) -> &mut Self {
         self.args(["--memory", self.guest.default_memory_string().as_str()])
     }