allow plugins to take namespace and nargs

mplsgrant · mplsgrant · commit 691cff5d9b6b · 2024-12-23T00:08:15.000-06:00
Also, add a hello-world test
diff --git a/resources/plugins/simln/plugin.py b/resources/plugins/simln/plugin.py
@@ -10,7 +10,7 @@
 
 # When we want to select pods based on their role in Warnet, we use "mission" tags. The "mission"
 # tag for "lightning" nodes is stored in LIGHTNING_MISSION.
-from warnet.constants import LIGHTNING_MISSION
+from warnet.constants import LIGHTNING_MISSION, HookValue
 from warnet.k8s import (
     download,
     get_default_namespace,
@@ -56,13 +56,24 @@ def simln(ctx):
     ctx.obj[PLUGIN_DIR_TAG] = Path(plugin_dir)
 
 
+# Each Warnet plugin must have an entrypoint function which takes a network_file_path and a
+# hook_value. Possible hook values can be found in the HookValue enum. It also takes a namespace
+# value and a variable number of arguments which is used by, for example, preNode and postNode to
+# pass along node names.
 @simln.command()
 @click.argument("network_file_path", type=str)
 @click.argument("hook_value", type=str)
+@click.argument("namespace", type=str)
+@click.argument("nargs", nargs=-1)
 @click.pass_context
-def entrypoint(ctx, network_file_path: str, hook_value: str):
+def entrypoint(ctx, network_file_path: str, hook_value: str, namespace: str, nargs):
     """Plugin entrypoint"""
+    assert hook_value in {
+        item.value for item in HookValue
+    }, f"{hook_value} is not a valid HookValue"
+
     network_file_path = Path(network_file_path)
+
     with network_file_path.open() as f:
         network_file = yaml.safe_load(f) or {}
         if not isinstance(network_file, dict):
@@ -76,11 +87,11 @@ def entrypoint(ctx, network_file_path: str, hook_value: str):
     if not plugin_data:
         raise PluginError(f"Could not find {plugin_name} in {network_file_path}")
 
-    _entrypoint(ctx, plugin_data)
+    _entrypoint(ctx, plugin_data, HookValue(hook_value), namespace, nargs)
 
 
-def _entrypoint(ctx, plugin_data: dict):
-    """ "Called by entrypoint"""
+def _entrypoint(ctx, plugin_data: dict, hook_value: HookValue, namespace: str, nargs):
+    """Called by entrypoint"""
     # write your plugin startup commands here
     activity = plugin_data.get("activity")
     activity = json.loads(activity)
diff --git a/src/warnet/deploy.py b/src/warnet/deploy.py
@@ -85,7 +85,7 @@ def _deploy(directory, debug, namespace, to_all_users):
         return
 
     if (directory / NETWORK_FILE).exists():
-        run_plugins(directory, HookValue.PRE_DEPLOY)
+        run_plugins(directory, HookValue.PRE_DEPLOY, namespace)
 
         processes = []
         # Deploy logging CRD first to avoid synchronisation issues
@@ -95,7 +95,7 @@ def _deploy(directory, debug, namespace, to_all_users):
         logging_process.start()
         processes.append(logging_process)
 
-        run_plugins(directory, HookValue.PRE_NETWORK)
+        run_plugins(directory, HookValue.PRE_NETWORK, namespace)
 
         network_process = Process(target=deploy_network, args=(directory, debug, namespace))
         network_process.start()
@@ -111,7 +111,7 @@ def _deploy(directory, debug, namespace, to_all_users):
         # Wait for the network process to complete
         network_process.join()
 
-        run_plugins(directory, HookValue.POST_NETWORK)
+        run_plugins(directory, HookValue.POST_NETWORK, namespace)
 
         # Start the fork observer process immediately after network process completes
         fork_observer_process = Process(target=deploy_fork_observer, args=(directory, debug))
@@ -122,7 +122,7 @@ def _deploy(directory, debug, namespace, to_all_users):
         for p in processes:
             p.join()
 
-        run_plugins(directory, HookValue.POST_DEPLOY)
+        run_plugins(directory, HookValue.POST_DEPLOY, namespace)
 
     elif (directory / NAMESPACES_FILE).exists():
         deploy_namespaces(directory)
@@ -132,7 +132,7 @@ def _deploy(directory, debug, namespace, to_all_users):
         )
 
 
-def run_plugins(directory, hook_value: HookValue):
+def run_plugins(directory, hook_value: HookValue, namespace, *args):
     """Run the plugin commands within a given hook value"""
 
     network_file_path = directory / NETWORK_FILE
@@ -154,7 +154,7 @@ def run_plugins(directory, hook_value: HookValue):
                 except Exception as err:
                     raise SyntaxError("Each plugin must have an 'entrypoint'") from err
 
-                cmd = f"{network_file_path.parent / entrypoint_path / Path('plugin.py')} entrypoint {network_file_path} {hook_value.value}"
+                cmd = f"{network_file_path.parent / entrypoint_path / Path('plugin.py')} entrypoint {network_file_path} {hook_value.value} {namespace} {' '.join(map(str, args))}"
                 process = Process(target=run_command, args=(cmd,))
                 processes.append(process)
 
@@ -385,13 +385,13 @@ def deploy_single_node(node, directory: Path, debug: bool, namespace: str):
                 temp_override_file_path = Path(temp_file.name)
             cmd = f"{cmd} -f {temp_override_file_path}"
 
-        run_plugins(directory, HookValue.PRE_NODE)
+        run_plugins(directory, HookValue.PRE_NODE, namespace, node_name)
 
         if not stream_command(cmd):
             click.echo(f"Failed to run Helm command: {cmd}")
             return
 
-        run_plugins(directory, HookValue.POST_NODE)
+        run_plugins(directory, HookValue.POST_NODE, namespace, node_name)
 
     except Exception as e:
         click.echo(f"Error: {e}")
diff --git a/test/data/ln/network.yaml b/test/data/ln/network.yaml
@@ -54,8 +54,31 @@ nodes:
       lnd: true
 
 plugins:
+  preDeploy:
+    hello:
+      entrypoint: "../plugins/hello"
+      podName: "hello-pre-deploy"
+      helloTo: "preDeploy!"
   postDeploy:
     simln:
       entrypoint: "../../../resources/plugins/simln"
       activity: '[{"source": "tank-0003-ln", "destination": "tank-0005-ln", "interval_secs": 1, "amount_msat": 2000}]'
+  preNode:
+    hello:
+      entrypoint: "../plugins/hello"
+      helloTo: "preNode!"
+  postNode:
+    hello:
+      entrypoint: "../plugins/hello"
+      helloTo: "preNode!"
+  preNetwork:
+    hello:
+      entrypoint: "../plugins/hello"
+      helloTo: "preNetwork!"
+      podName: "hello-pre-network"
+  postNetwork:
+    hello:
+      entrypoint: "../plugins/hello"
+      helloTo: "postNetwork!"
+      podName: "hello-post-network"
 
diff --git a/test/data/plugins/hello/charts/hello/Chart.yaml b/test/data/plugins/hello/charts/hello/Chart.yaml
@@ -0,0 +1,5 @@
+apiVersion: v2
+name: hello-chart
+description: A Helm chart for a hello Pod
+version: 0.1.0
+appVersion: "1.0"
diff --git a/test/data/plugins/hello/charts/hello/templates/pod.yaml b/test/data/plugins/hello/charts/hello/templates/pod.yaml
@@ -0,0 +1,15 @@
+apiVersion: v1
+kind: Pod
+metadata:
+  name: {{ .Values.podName }}
+  labels:
+    app: {{ .Chart.Name }}
+spec:
+  restartPolicy: Never
+  containers:
+  - name: {{ .Values.podName }}-container
+    image: alpine:latest
+    command: ["sh", "-c"]
+    args:
+      - echo "Hello {{ .Values.helloTo }}";
+    resources: {}
diff --git a/test/data/plugins/hello/charts/hello/values.yaml b/test/data/plugins/hello/charts/hello/values.yaml
@@ -0,0 +1,2 @@
+podName: hello-pod
+helloTo: "world"
diff --git a/test/data/plugins/hello/plugin.py b/test/data/plugins/hello/plugin.py
@@ -0,0 +1,159 @@
+#!/usr/bin/env python3
+import logging
+from pathlib import Path
+from typing import Optional
+
+import click
+import yaml
+from kubernetes.stream import stream
+
+from warnet.constants import HookValue
+from warnet.k8s import (
+    get_default_namespace,
+    get_static_client,
+)
+from warnet.process import run_command
+
+MISSION = "hello"
+PRIMARY_CONTAINER = MISSION
+
+PLUGIN_DIR_TAG = "plugin_dir"
+
+
+class PluginError(Exception):
+    pass
+
+
+log = logging.getLogger(MISSION)
+if not log.hasHandlers():
+    console_handler = logging.StreamHandler()
+    console_handler.setLevel(logging.DEBUG)
+    formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+    console_handler.setFormatter(formatter)
+    log.addHandler(console_handler)
+log.setLevel(logging.DEBUG)
+log.propagate = True
+
+
+@click.group()
+@click.pass_context
+def hello(ctx):
+    """Commands for the Hello plugin"""
+    ctx.ensure_object(dict)
+    plugin_dir = Path(__file__).resolve().parent
+    ctx.obj[PLUGIN_DIR_TAG] = Path(plugin_dir)
+
+
+@hello.command()
+@click.argument("network_file_path", type=str)
+@click.argument("hook_value", type=str)
+@click.argument("namespace", type=str)
+@click.argument("nargs", nargs=-1)
+@click.pass_context
+def entrypoint(ctx, network_file_path: str, hook_value: str, namespace: str, nargs):
+    """Plugin entrypoint"""
+    assert hook_value in {
+        item.value for item in HookValue
+    }, f"{hook_value} is not a valid HookValue"
+
+    network_file_path = Path(network_file_path)
+
+    with network_file_path.open() as f:
+        network_file = yaml.safe_load(f) or {}
+        if not isinstance(network_file, dict):
+            raise ValueError(f"Invalid network file structure: {network_file_path}")
+
+    plugins_section = network_file.get("plugins", {})
+    hook_section = plugins_section.get(hook_value, {})
+
+    plugin_name = Path(__file__).resolve().parent.stem
+    plugin_data = hook_section.get(plugin_name)
+    if not plugin_data:
+        raise PluginError(f"Could not find {plugin_name} in {network_file_path}")
+
+    _entrypoint(ctx, plugin_data, HookValue(hook_value), namespace, nargs)
+
+
+def _entrypoint(ctx, plugin_data: dict, hook_value: HookValue, namespace: str, nargs):
+    """Called by entrypoint"""
+    match hook_value:
+        case (
+            HookValue.PRE_NETWORK
+            | HookValue.POST_NETWORK
+            | HookValue.PRE_DEPLOY
+            | HookValue.POST_DEPLOY
+        ):
+            data = get_data(plugin_data)
+            if data:
+                _launch_pod(ctx, install_name=hook_value.value.lower() + "-hello", **data)
+            else:
+                _launch_pod(ctx, install_name=hook_value.value.lower() + "-hello")
+        case HookValue.PRE_NODE:
+            name = nargs[0] + "-pre-hello-pod"
+            _launch_pod(ctx, install_name=hook_value.value.lower() + "-" + name, podName=name)
+        case HookValue.POST_NODE:
+            name = nargs[0] + "-post-hello-pod"
+            _launch_pod(ctx, install_name=hook_value.value.lower() + "-" + name, podName=name)
+
+
+def get_data(plugin_data: dict) -> Optional[dict]:
+    data = {key: plugin_data.get(key) for key in ("podName", "helloTo") if plugin_data.get(key)}
+    return data or None
+
+
+def _launch_pod(
+    ctx, install_name: str = "hello", podName: str = "hello-pod", helloTo: str = "World!"
+):
+    command = f"helm upgrade --install {install_name} {ctx.obj[PLUGIN_DIR_TAG]}/charts/hello --set podName={podName} --set helloTo={helloTo}"
+    log.info(command)
+    log.info(run_command(command))
+
+
+def _sh(pod, method: str, params: tuple[str, ...]) -> str:
+    namespace = get_default_namespace()
+
+    sclient = get_static_client()
+    if params:
+        cmd = [method]
+        cmd.extend(params)
+    else:
+        cmd = [method]
+    try:
+        resp = stream(
+            sclient.connect_get_namespaced_pod_exec,
+            pod,
+            namespace,
+            container=PRIMARY_CONTAINER,
+            command=cmd,
+            stderr=True,
+            stdin=False,
+            stdout=True,
+            tty=False,
+            _preload_content=False,
+        )
+        stdout = ""
+        stderr = ""
+        while resp.is_open():
+            resp.update(timeout=1)
+            if resp.peek_stdout():
+                stdout_chunk = resp.read_stdout()
+                stdout += stdout_chunk
+            if resp.peek_stderr():
+                stderr_chunk = resp.read_stderr()
+                stderr += stderr_chunk
+        return stdout + stderr
+    except Exception as err:
+        print(f"Could not execute stream: {err}")
+
+
+@hello.command(context_settings={"ignore_unknown_options": True})
+@click.argument("pod", type=str)
+@click.argument("method", type=str)
+@click.argument("params", type=str, nargs=-1)  # this will capture all remaining arguments
+def sh(pod: str, method: str, params: tuple[str, ...]):
+    """Run shell commands in a pod"""
+    print(_sh(pod, method, params))
+
+
+if __name__ == "__main__":
+    hello()
diff --git a/test/simln_test.py b/test/simln_test.py
@@ -25,6 +25,7 @@ def run_test(self):
             self.init_directory()
             self.deploy_with_plugin()
             self.copy_results()
+            self.assert_hello_plugin()
         finally:
             self.cleanup()
 
@@ -83,6 +84,24 @@ def found_results_locally(self) -> bool:
         self.log.info(f"Did not find downloaded results in directory: {directory}.")
         return False
 
+    def assert_hello_plugin(self):
+        self.log.info("Waiting for the 'hello' plugin pods.")
+        wait_for_pod("hello-pre-deploy")  # We don't use post-deploy (simln covers that)
+        wait_for_pod("hello-pre-network")
+        wait_for_pod("hello-post-network")
+        wait_for_pod("tank-0000-post-hello-pod")
+        wait_for_pod("tank-0000-pre-hello-pod")
+        wait_for_pod("tank-0001-post-hello-pod")
+        wait_for_pod("tank-0001-pre-hello-pod")
+        wait_for_pod("tank-0002-post-hello-pod")
+        wait_for_pod("tank-0002-pre-hello-pod")
+        wait_for_pod("tank-0003-post-hello-pod")
+        wait_for_pod("tank-0003-pre-hello-pod")
+        wait_for_pod("tank-0004-post-hello-pod")
+        wait_for_pod("tank-0004-pre-hello-pod")
+        wait_for_pod("tank-0005-post-hello-pod")
+        wait_for_pod("tank-0005-pre-hello-pod")
+
 
 if __name__ == "__main__":
     test = SimLNTest()

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,2 @@`
	`1`	`+podName: hello-pod`
	`2`	`+helloTo: "world"`