From 85159c9255b35c990fb636e68cd4b7d4ca953b84 Mon Sep 17 00:00:00 2001
From: Muminul Islam <muislam@microsoft.com>
Date: Fri, 27 Feb 2026 18:17:40 -0800
Subject: [PATCH] tests: Add factory pattern to create guest instances

Introduce GuestFactory struct that encapsulates GuestVmType,
boot timeout, and nested virtualization defaults. This avoids
repeatedly specifying VM type and associated defaults when
constructing Guest instances in integration tests.

Two factory constructors are provided:
- new_regular_guest_factory: default timeout, nested enabled
- new_confidential_guest_factory: CVM timeout, nested disabled

Multiple create_guest variants allow customizing CPU count,
memory size, and nested virtualization while inheriting the
factory's VM type and timeout settings.

Signed-off-by: Muminul Islam <muislam@microsoft.com>
---
 test_infra/src/lib.rs | 54 +++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 54 insertions(+)

diff --git a/test_infra/src/lib.rs b/test_infra/src/lib.rs
index 345aa3c25..1ba5ceccc 100644
--- a/test_infra/src/lib.rs
+++ b/test_infra/src/lib.rs
@@ -998,6 +998,21 @@ impl Guest {
         Self::new_from_ip_range(disk_config, "192.168", next_guest_id())
     }
 
+    pub fn with_cpu(mut self, count: u32) -> Self {
+        self.num_cpu = count;
+        self
+    }
+
+    pub fn with_memory(mut self, mem_size: &str) -> Self {
+        self.mem_size_str = mem_size.to_string();
+        self
+    }
+
+    pub fn with_nested(mut self, nested: bool) -> Self {
+        self.nested = nested;
+        self
+    }
+
     pub fn default_net_string(&self) -> String {
         format!(
             "tap=,mac={},ip={},mask=255.255.255.128",
@@ -1435,6 +1450,45 @@ impl Guest {
     }
 }
 
+// A factory for creating guests with different configurations. The factory is initialized
+// with a GuestVmType, and created guests will have the same GuestVmType as the factory.
+// This allows creation of guests with different configurations (e.g. regular vs confidential)
+// without specifying the GuestVmType each time.
+// Based on the VmType, the default timeout for waiting for the VM to boot is also set,
+// which is used in the wait_vm_boot() method of the Guest struct. Additionally, nested
+// virtualization is disabled by default for confidential VMs, as it is not supported.
+pub struct GuestFactory {
+    vm_type: GuestVmType,
+    boot_timeout: u32,
+    nested: bool,
+}
+
+impl GuestFactory {
+    pub fn new_regular_guest_factory() -> Self {
+        Self {
+            vm_type: GuestVmType::Regular,
+            boot_timeout: DEFAULT_TCP_LISTENER_TIMEOUT,
+            nested: true,
+        }
+    }
+
+    pub fn new_confidential_guest_factory() -> Self {
+        Self {
+            vm_type: GuestVmType::Confidential,
+            boot_timeout: DEFAULT_CVM_TCP_LISTENER_TIMEOUT,
+            nested: false,
+        }
+    }
+
+    pub fn create_guest(&self, disk_config: Box<dyn DiskConfig>) -> Guest {
+        let mut guest = Guest::new(disk_config);
+        guest.vm_type = self.vm_type;
+        guest.boot_timeout = self.boot_timeout;
+        guest.nested = self.nested;
+        guest
+    }
+}
+
 #[derive(Default)]
 pub enum VerbosityLevel {
     #[default]