Add CPU and memory config options to test framework

4 files changed, 53 insertions, 2 deletions

diff --git a/test/test-manager/src/config/vm.rs b/test/test-manager/src/config/vm.rs
index 911d2fcf64..144c72b136 100644
--- a/test/test-manager/src/config/vm.rs
+++ b/test/test-manager/src/config/vm.rs
@@ -60,6 +60,16 @@ pub struct VmConfig {
     #[serde(default)]
     #[arg(long, requires("tpm"))]
     pub ovmf_code_path: Option<String>,
+
+    /// Number of vCPUs
+    #[serde(default)]
+    #[arg(long)]
+    pub vcpus: Option<usize>,
+
+    /// Amount of memory, in MBs
+    #[serde(default)]
+    #[arg(long)]
+    pub memory: Option<usize>,
 }
 
 impl VmConfig {
diff --git a/test/test-manager/src/main.rs b/test/test-manager/src/main.rs
index 694af65dc3..2a043b4276 100644
--- a/test/test-manager/src/main.rs
+++ b/test/test-manager/src/main.rs
@@ -31,6 +31,7 @@ struct Args {
 }
 
 #[derive(clap::Subcommand, Debug)]
+#[allow(clippy::large_enum_variant)]
 enum Commands {
     /// Manage configuration for tests and VMs
     #[clap(subcommand)]
@@ -148,6 +149,7 @@ enum Commands {
 }
 
 #[derive(clap::Subcommand, Debug)]
+#[allow(clippy::large_enum_variant)]
 enum ConfigArg {
     /// Print the current config
     Get,
@@ -159,6 +161,7 @@ enum ConfigArg {
 }
 
 #[derive(clap::Subcommand, Debug)]
+#[allow(clippy::large_enum_variant)]
 enum VmConfig {
     /// Create or edit a VM config
     Set {
diff --git a/test/test-manager/src/vm/qemu.rs b/test/test-manager/src/vm/qemu.rs
index 3d6921e97d..9028bb94d9 100644
--- a/test/test-manager/src/vm/qemu.rs
+++ b/test/test-manager/src/vm/qemu.rs
@@ -25,6 +25,12 @@ const STDERR_LOG_LEVEL: log::Level = log::Level::Error;
 const STDOUT_LOG_LEVEL: log::Level = log::Level::Debug;
 const OBTAIN_IP_TIMEOUT: Duration = Duration::from_secs(180);
 
+/// Default number of VCPU cores (passed to -smp)
+const DEFAULT_NUM_VCPUS: usize = 2;
+
+/// Default amount of memory, in MBs (passed to -m)
+const DEFAULT_AMOUNT_MEMORY: usize = 4096;
+
 #[derive(thiserror::Error, Debug)]
 pub enum Error {
     #[error("Failed to set up network")]
@@ -83,15 +89,21 @@ pub async fn run(config: &Config, vm_config: &VmConfig) -> Result<QemuInstance>
         .map_err(Error::Network)?;
 
     let mut qemu_cmd = Command::new("qemu-system-x86_64");
+    let vcpus = vm_config.vcpus.unwrap_or(DEFAULT_NUM_VCPUS);
+    let memory = vm_config.memory.unwrap_or(DEFAULT_AMOUNT_MEMORY);
+
+    log::debug!("CPU count: {vcpus}");
+    log::debug!("Memory: {memory}M");
+
     qemu_cmd.args([
         "-cpu",
         "host",
         "-accel",
         "kvm",
         "-m",
-        "4096",
+        &memory.to_string(),
         "-smp",
-        "2",
+        &vcpus.to_string(),
         "-drive",
         &format!("file={}", vm_config.image_path),
         "-device",
diff --git a/test/test-manager/src/vm/tart.rs b/test/test-manager/src/vm/tart.rs
index 407a4903c2..faf6853d90 100644
--- a/test/test-manager/src/vm/tart.rs
+++ b/test/test-manager/src/vm/tart.rs
@@ -49,6 +49,10 @@ pub async fn run(config: &Config, vm_config: &VmConfig) -> Result<TartInstance>
         MachineCopy::clone_vm(&vm_config.image_path).await?
     };
 
+    if let Err(err) = machine_copy.configure(vm_config).await {
+        log::error!("Failed to configure tart vm: {err}");
+    }
+
     // Start VM
     let mut tart_cmd = Command::new("tart");
     tart_cmd.args(["run", &machine_copy.name, "--serial"]);
@@ -167,6 +171,28 @@ impl MachineCopy {
         })
     }
 
+    pub async fn configure(&self, vm_config: &VmConfig) -> Result<()> {
+        let mut args = vec![];
+        if let Some(cpu) = vm_config.vcpus {
+            args.extend(["--cpu".to_owned(), cpu.to_string()]);
+            log::info!("vCPUs: {cpu}");
+        }
+        if let Some(mem) = vm_config.memory {
+            args.extend(["--memory".to_owned(), mem.to_string()]);
+            log::info!("Memory: {mem} MB");
+        }
+        if !args.is_empty() {
+            let mut tart_cmd = Command::new("tart");
+            tart_cmd.args(["set", &self.name]);
+            tart_cmd.args(args);
+            tart_cmd
+                .status()
+                .await
+                .context("failed to update tart config")?;
+        }
+        Ok(())
+    }
+
     pub async fn cleanup(mut self) {
         let _ = tokio::task::spawn_blocking(move || self.try_destroy()).await;
     }