control: add --debug option to run WIP scenarios

Author: pinheadmz

src/warnet/control.py

@@ -17,11 +17,13 @@
 
 from .constants import COMMANDER_CHART, LOGGING_NAMESPACE
 from .k8s import (
+    delete_pod,
     get_default_namespace,
     get_mission,
     get_pods,
+    pod_log,
     snapshot_bitcoin_datadir,
-    pod_log
+    wait_for_pod,
 )
 from .process import run_command, stream_command
 
@@ -161,8 +163,14 @@ def get_active_network(namespace):
 
 @click.command(context_settings={"ignore_unknown_options": True})
 @click.argument("scenario_file", type=click.Path(exists=True, file_okay=True, dir_okay=False))
+@click.option(
+    "--debug",
+    is_flag=True,
+    default=False,
+    help="Stream scenario output and delete container when stopped",
+)
 @click.argument("additional_args", nargs=-1, type=click.UNPROCESSED)
-def run(scenario_file: str, additional_args: tuple[str]):
+def run(scenario_file: str, debug: bool, additional_args: tuple[str]):
     """
     Run a scenario from a file.
     Pass `-- --help` to get individual scenario help
@@ -230,11 +238,22 @@ def run(scenario_file: str, additional_args: tuple[str]):
         print(f"Failed to start scenario: {scenario_name}")
         print(f"Error: {e.stderr}")
 
+    if debug:
+        print("Waiting for commander pod to start...")
+        wait_for_pod(name)
+        _logs(pod_name=name, follow=True)
+        print("Deleting pod...")
+        delete_pod(name)
+
 
 @click.command()
 @click.argument("pod_name", type=str, default="")
 @click.option("--follow", "-f", is_flag=True, default=False, help="Follow logs")
 def logs(pod_name: str, follow: bool):
+    return _logs(pod_name, follow)
+
+
+def _logs(pod_name: str, follow: bool):
     """Show the logs of a pod"""
     namespace = get_default_namespace()
 
@@ -266,9 +285,11 @@ def logs(pod_name: str, follow: bool):
     try:
         stream = pod_log(pod_name, container_name=None, follow=follow)
         for line in stream.stream():
-            print(line.decode('utf-8'), end=None)
+            print(line.decode("utf-8"), end=None)
     except Exception as e:
         print(e)
+    except KeyboardInterrupt:
+        print("Interrupted streaming log!")
 
 
 @click.command()

src/warnet/k8s.py

@@ -3,13 +3,14 @@
 import sys
 import tempfile
 from pathlib import Path
+from time import sleep
 
 import yaml
 from kubernetes import client, config, watch
 from kubernetes.client.models import CoreV1Event, V1PodList
+from kubernetes.client.rest import ApiException
 from kubernetes.dynamic import DynamicClient
 from kubernetes.stream import stream
-from kubernetes.client.rest import ApiException
 
 from .constants import (
     CADDY_INGRESS_NAME,
@@ -293,7 +294,17 @@ def pod_log(pod_name, container_name=None, follow=False):
             namespace=get_default_namespace(),
             container=container_name,
             follow=follow,
-            _preload_content=False
+            _preload_content=False,
         )
     except ApiException as e:
-        raise Exception(json.loads(e.body.decode('utf-8'))["message"])
+        raise Exception(json.loads(e.body.decode("utf-8"))["message"]) from None
+
+
+def wait_for_pod(pod_name, timeout_seconds=10):
+    sclient = get_static_client()
+    while timeout_seconds > 0:
+        pod = sclient.read_namespaced_pod_status(name=pod_name, namespace=get_default_namespace())
+        if pod.status.phase != "Pending":
+            return
+        sleep(1)
+        timeout_seconds -= 1