From 7abc67c013621509461241528c2f5e7de00f05a7 Mon Sep 17 00:00:00 2001
From: Kelvin Zhang <zhangkelvin@google.com>
Date: Fri, 21 Jun 2024 11:05:13 -0700
Subject: [PATCH] Add uefi_load unittest to github CI

---
 .github/workflows/github-ci-clang.yml | 12 ++++
 unittest.py                           | 85 +++++++++++++++++++++++++++
 2 files changed, 97 insertions(+)
 create mode 100644 unittest.py

diff --git a/.github/workflows/github-ci-clang.yml b/.github/workflows/github-ci-clang.yml
index 9c036eea6..0811f7e59 100644
--- a/.github/workflows/github-ci-clang.yml
+++ b/.github/workflows/github-ci-clang.yml
@@ -73,5 +73,17 @@ jobs:
       shell: bash
       run: |
         make -j $(nproc)
+    # When LK is compiled with DEBUG=0, there's no console and no way for us
+    # to read test output
+    - name: qemu
+      if: ${{ matrix.project == 'qemu-virt-arm64-test' && matrix.debug == 2 }}
+      shell: bash
+      run: |
+        env -i DEBIAN_FRONTEND=noninteractive sudo apt install -y qemu-system-arm
+    - name: test uefi_load
+      if: ${{ matrix.project == 'qemu-virt-arm64-test' && matrix.debug == 2 }}
+      shell: bash
+      run: |
+        python3 unittest.py
 
 # vim: ts=2 sw=2 expandtab
diff --git a/unittest.py b/unittest.py
new file mode 100644
index 000000000..0bad57854
--- /dev/null
+++ b/unittest.py
@@ -0,0 +1,85 @@
+import os
+import sys
+import signal
+import subprocess
+from subprocess import PIPE, STDOUT
+import io
+import time
+import select
+
+
+def wait_for_output(fp: io.TextIOBase, predicate, timeout: float):
+    start_time = time.time()
+    end_time = start_time + timeout
+    poll_obj = select.poll()
+    poll_obj.register(fp, select.POLLIN)
+    output = []
+    cur_line = ""
+    while time.time() < end_time:
+        poll_result = poll_obj.poll(max(end_time-time.time(), 0.001))
+        if poll_result:
+            data = fp.read()
+            if "\n" in data:
+                output.append(cur_line + data[:data.index("\n") + 1])
+                cur_line = data[data.index("\n") + 1:]
+                if predicate(output[-1]):
+                    return True, output
+            else:
+                cur_line += data
+
+    return False, output
+
+
+def shutdown_little_kernel(p: subprocess.Popen):
+    try:
+        p.stdin.write("poweroff\n")
+        p.stdin.flush()
+        p.wait(0.3)
+        p.send_signal(signal.SIGINT)
+        p.wait(1)
+    except subprocess.TimeoutExpired:
+        pass
+    finally:
+        p.kill()
+        p.wait()
+
+
+def main():
+    p = subprocess.Popen(['qemu-system-aarch64',
+                          '-cpu',
+                          'max',
+                          '-m',
+                          '512',
+                          '-smp',
+                          '1',
+                          '-machine',
+                          'virt,highmem=off',
+                          '-kernel',
+                          'build-qemu-virt-arm64-test/lk.elf',
+                          '-net',
+                          'none',
+                          '-nographic',
+                          '-drive',
+                          'if=none,file=lib/uefi/helloworld_aa64.efi,id=blk,format=raw',
+                          '-device',
+                          'virtio-blk-device,drive=blk'], stdout=PIPE, stdin=PIPE, stderr=STDOUT, text=True)
+    try:
+        os.set_blocking(p.stdout.fileno(), False)
+        condition_met, output = wait_for_output(
+            p.stdout, lambda l: "starting app shell" in l, 0.5)
+        assert condition_met, "Did not see 'starting app shell', stdout: {}".format(
+            "".join(output))
+        p.stdin.write("uefi_load virtio0\n")
+        p.stdin.flush()
+        condition_met, output = wait_for_output(
+            p.stdout, lambda l: "Hello World!" in l, 0.5)
+        print("".join(output))
+        if condition_met:
+            return
+
+    finally:
+        shutdown_little_kernel(p)
+
+
+if __name__ == "__main__":
+    main()