From 3e258a7f97d41a829887af355577f4e5e8295706 Mon Sep 17 00:00:00 2001
From: Ramyak mehra <rmehra_be19@thapar.edu>
Date: Fri, 29 Jul 2022 20:30:32 +0000
Subject: [PATCH] ported x86_64 python tests to be compatible with aarch64

Added support to pull and set aach64 images
Changed how vcpus are checked. This is done because on
aarch64 /proc/cpuinfo doesn't directly return the number of
cpus.

Signed-off-by: Ramyak mehra <rmehra_be19@thapar.edu>
---
 .buildkite/custom-tests.json    |  3 +-
 tests/test_run_reference_vmm.py | 68 ++++++++++++++++++++++++---------
 2 files changed, 51 insertions(+), 20 deletions(-)

diff --git a/.buildkite/custom-tests.json b/.buildkite/custom-tests.json
index 70c5ff6f..90e840d3 100644
--- a/.buildkite/custom-tests.json
+++ b/.buildkite/custom-tests.json
@@ -4,7 +4,8 @@
             "test_name": "run",
             "command": "pytest tests/test_run_reference_vmm.py",
             "platform": [
-                "x86_64"
+                "x86_64",
+                "aarch64"
             ]
         }
     ]
diff --git a/tests/test_run_reference_vmm.py b/tests/test_run_reference_vmm.py
index dd0a96cc..0ce36f8d 100644
--- a/tests/test_run_reference_vmm.py
+++ b/tests/test_run_reference_vmm.py
@@ -6,6 +6,7 @@
 import os
 import subprocess
 import tempfile
+import platform
 
 import time
 from subprocess import PIPE
@@ -17,6 +18,7 @@
 # No. of seconds after which to give up for the test
 TEST_TIMEOUT = 30
 
+arch = platform.machine()
 
 def process_exists(pid):
     try:
@@ -58,24 +60,43 @@ def default_ubuntu_elf():
         first=True
     )
 
+def default_ubuntu_pe():
+    return s3_download(
+        resource_name="pe-focal",
+        resource_type="kernel",
+        first=True
+    )
+
+def default_busybox_pe():
+    return s3_download(
+        resource_name="pe-hello-busybox",
+        resource_type="kernel",
+        first=True
+    )
 
 def default_disk():
     return s3_download(
-        resource_name="ubuntu-focal-rootfs-x86_64.ext4",
+        resource_name=f"ubuntu-focal-rootfs-{arch}.ext4",
         resource_type="disk",
         first=True
     )
 
+if arch == "aarch64":
+    KERNELS_INITRAMFS = [default_busybox_pe()]
 
-KERNELS_INITRAMFS = [
-    default_busybox_elf(),
-    default_busybox_bzimage()
-]
+    UBUNTU_KERNEL_DISK_PAIRS = [
+        (default_ubuntu_pe() , default_disk())
+    ]
+else:
+    KERNELS_INITRAMFS = [
+        default_busybox_elf(),
+        default_busybox_bzimage()
+    ]
 
-UBUNTU_KERNEL_DISK_PAIRS = [
-    (default_ubuntu_elf(), default_disk()),
-    (default_ubuntu_bzimage(), default_disk()),
-]
+    UBUNTU_KERNEL_DISK_PAIRS = [
+        (default_ubuntu_elf(), default_disk()),
+        (default_ubuntu_bzimage(), default_disk()),
+    ]
 
 
 """
@@ -101,13 +122,16 @@ def start_vmm_process(kernel_path, disk_path=None, num_vcpus=1, mem_size_mib=102
         "--kernel"
     ]
     if default_cmdline:
-        kernel_config = "path={},kernel_load_addr={}".format(
-            kernel_path, kernel_load_addr
+        kernel_config = "path={}".format(
+            kernel_path
         )
     else:
-        kernel_config = "cmdline={},path={},kernel_load_addr={}".format(
-            cmdline, kernel_path, kernel_load_addr
+        kernel_config = "cmdline={},path={}".format(
+            cmdline, kernel_path
         )
+    if arch == "x86_64":
+        kernel_config += ",kernel_load_addr={}".format(kernel_load_addr)
+    
     vmm_cmd.append(kernel_config)
     
     tmp_file_path = None
@@ -232,13 +256,19 @@ def expect_vcpus(vmm_process, expected_vcpus):
     # dependent on it, so it's just fine to call it again, less than ideal, but not
     # wrong.
     setup_stdout_nonblocking(vmm_process)
-
+    # We check for the prompt at the start so that the inital stdout of the process
+    # gets cleared up, which is printed while the machine is booted.
+    prompt = '/ #'
+    expect_string(vmm_process, prompt)
+    
     # /proc/cpuinfo displays info about each vCPU
-    vmm_process.stdin.write(b'cat /proc/cpuinfo\n')
-    vmm_process.stdin.flush()
-
-    siblings_line = expect_string(vmm_process, "siblings")
-    actual_vcpus = int(siblings_line.split(":")[1].strip())
+    cmd = 'cat /proc/cpuinfo'
+    output = run_cmd_inside_vm(cmd.encode(), vmm_process, prompt.encode(), timeout=5)
+    
+    actual_vcpus = 0
+    for line in output:
+        if "processor" in line.decode():
+            actual_vcpus +=1
 
     assert actual_vcpus == expected_vcpus, \
         "Expected {}, found {} vCPUs".format(expected_vcpus, actual_vcpus)