From 879919b1093fb51ee8f53439ecdd13357d7c4ac8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Pablo=20Barb=C3=A1chano?= <pablob@amazon.com>
Date: Fri, 15 Sep 2023 13:55:28 +0200
Subject: [PATCH] test: add sandbox tool
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Add a sandbox tool to run experiments and to do quick checks.

Signed-off-by: Pablo Barbáchano <pablob@amazon.com>
---
 tests/README.md  | 17 +++++++++++-
 tools/devtool    |  7 +++++
 tools/sandbox.py | 67 ++++++++++++++++++++++++++++++++++++++++++++++++
 tools/test.sh    |  1 +
 4 files changed, 91 insertions(+), 1 deletion(-)
 create mode 100755 tools/sandbox.py

diff --git a/tests/README.md b/tests/README.md
index 2880785fffc4..c3a3aca9e65b 100644
--- a/tests/README.md
+++ b/tests/README.md
@@ -306,7 +306,6 @@ can examine local variables and the stack, and can use the normal Python REPL.
 
 ```sh
 ./tools/devtool -y shell --privileged
-pip3 install ipython
 export PYTEST_ADDOPTS=--pdbcls=IPython.terminal.debugger:TerminalPdb
 ./tools/test.sh -k 1024 integration_tests/performance/test_boottime.py::test_boottime
 ```
@@ -378,3 +377,19 @@ ipdb> test_microvm.help.gdbserver()
 ```
 
 You get some instructions on how to run GDB to attach to gdbserver.
+
+## Sandbox
+
+```sh
+./tools/devtool sandbox
+```
+
+That should drop you in an IPython REPL, where you can interact with a microvm:
+
+```python
+uvm.help.print_log()
+uvm.get_all_metrics()
+uvm.ssh.run("ls")
+snap = uvm.snapshot_full()
+uvm.help.tmux_ssh()
+```
diff --git a/tools/devtool b/tools/devtool
index 1962e107ec98..9c724a346bff 100755
--- a/tools/devtool
+++ b/tools/devtool
@@ -414,6 +414,9 @@ cmd_help() {
     cat <<EOF
     test_debug [-- [<pytest args>]]
         Run tests in a debugging environment
+
+    test_sandbox
+        Run Firecracker in an IPython REPL
 EOF
 }
 
@@ -704,6 +707,10 @@ cmd_sh() {
         bash --norc -c "$*"
 }
 
+cmd_sandbox() {
+    cmd_sh "tmux new env PYTEST_ADDOPTS=--pdbcls=IPython.terminal.debugger:TerminalPdb PYTHONPATH=tests IPYTHONDIR=\$PWD/.ipython ipython -i ./tools/sandbox.py $@"
+}
+
 cmd_test_debug() {
     cmd_sh "tmux new ./tools/test.sh --pdb $@"
 }
diff --git a/tools/sandbox.py b/tools/sandbox.py
new file mode 100755
index 000000000000..655c55338878
--- /dev/null
+++ b/tools/sandbox.py
@@ -0,0 +1,67 @@
+#!/usr/bin/env python3
+# Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Run Firecracker in an IPython REPL
+"""
+
+import argparse
+import re
+
+from framework.artifacts import disks, kernels
+from framework.microvm import MicroVMFactory
+from host_tools.cargo_build import get_firecracker_binaries
+
+kernels = list(kernels("vmlinux-*"))
+rootfs = list(disks("ubuntu*ext4"))
+
+
+def parse_byte_size(param):
+    """
+    >>> parse_byte_size("1MB")
+    1048576
+    """
+    unit = {
+        "MB": 2**20,
+        "GB": 2**30,
+    }
+    match = re.match(r"(?P<val>\d+)(?P<unit>[MG]B)", param.upper())
+    return int(match.group("val")) * unit[match.group("unit")]
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    "--kernel",
+    choices=kernels,
+    default=kernels[-1],
+    help=f"Kernel to use. [{kernels[-1]}]",
+)
+parser.add_argument(
+    "--rootfs",
+    choices=rootfs,
+    default=rootfs[-1],
+    help=f"Rootfs to use. [{rootfs[-1]}]",
+)
+parser.add_argument("--vcpus", type=int, default=2)
+parser.add_argument(
+    "--guest-mem-size",
+    type=parse_byte_size,
+    default=128 * 2**20,  # 128MB
+)
+parser.add_argument("--rootfs-size", type=parse_byte_size, default=1 * 2**30)  # 1GB
+args = parser.parse_args()
+print(args)
+
+
+print("This step may take a while to compile Firecracker ...")
+vmfcty = MicroVMFactory("/srv", None, *get_firecracker_binaries())
+uvm = vmfcty.build(args.kernel, args.rootfs)
+uvm.help.enable_console()
+uvm.help.resize_disk(uvm.rootfs_file, args.rootfs_size)
+uvm.spawn()
+uvm.help.print_log()
+uvm.add_net_iface()
+uvm.basic_config(vcpu_count=args.vcpus, mem_size_mib=args.guest_mem_size // 2**20)
+uvm.start()
+uvm.get_all_metrics()
diff --git a/tools/test.sh b/tools/test.sh
index da0e24c29f3c..206ed4e9cd70 100755
--- a/tools/test.sh
+++ b/tools/test.sh
@@ -35,5 +35,6 @@ say "Copy CI artifacts to /srv, so hardlinks work"
 cp -ruvf build/img /srv
 
 cd tests
+export PYTEST_ADDOPTS="${PYTEST_ADDOPTS:-} --pdbcls=IPython.terminal.debugger:TerminalPdb"
 pytest "$@"
 exit $?