Make test_framework available for users when creating scenarios

pyproject.toml

@@ -51,7 +51,7 @@ build-backend = "setuptools.build_meta"
 include-package-data = true
 
 [tool.setuptools.packages.find]
-where = ["src", "."]
+where = ["src", ".", "resources/scenarios"]
 include = ["warnet*", "test_framework*", "resources*"]
 
 [tool.setuptools.package-data]

resources/charts/commander/templates/configmap.yaml

@@ -1,17 +1,17 @@
 apiVersion: v1
 kind: ConfigMap
 metadata:
-  name: {{ include "commander.fullname" . }}-scenario
+  name: {{ include "commander.fullname" . }}-warnet
   labels:
     {{- include "commander.labels" . | nindent 4 }}
 binaryData:
-  scenario.py: {{ .Values.scenario }}
+  warnet.json: {{ .Values.warnet }}
 ---
 apiVersion: v1
 kind: ConfigMap
 metadata:
-  name: {{ include "commander.fullname" . }}-warnet
+  name: {{ include "commander.fullname" . }}-archive
   labels:
     {{- include "commander.labels" . | nindent 4 }}
 binaryData:
-  warnet.json: {{ .Values.warnet }}
+  archive.pyz: {{ .Values.archive }}

resources/charts/commander/templates/pod.yaml

@@ -10,23 +10,23 @@ spec:
   restartPolicy: {{ .Values.restartPolicy }}
   containers:
     - name: {{ .Chart.Name }}
-      image: "{{ .Values.image.repository }}:{{ .Values.image.tag | default .Chart.AppVersion }}"
-      imagePullPolicy: {{ .Values.image.pullPolicy }}
+      image: python:3.12-slim
+      imagePullPolicy: IfNotPresent
       command: ["/bin/sh", "-c"]
       args:
         - |
-          python3 /scenario.py {{ .Values.args }}
+          python3 /archive.pyz {{ .Values.args }}
       volumeMounts:
-        - name: scenario
-          mountPath: /scenario.py
-          subPath: scenario.py
         - name: warnet
           mountPath: /warnet.json
           subPath: warnet.json
+        - name: archive
+          mountPath: /archive.pyz
+          subPath: archive.pyz
   volumes:
-    - name: scenario
-      configMap:
-        name: {{ include "commander.fullname" . }}-scenario
     - name: warnet
       configMap:
         name: {{ include "commander.fullname" . }}-warnet
+    - name: archive
+      configMap:
+        name: {{ include "commander.fullname" . }}-archive

resources/charts/commander/values.yaml

@@ -5,12 +5,6 @@ namespace: warnet
 
 restartPolicy: Never
 
-image:
-  repository: bitcoindevproject/warnet-commander
-  pullPolicy: IfNotPresent
-  # Overrides the image tag whose default is the chart appVersion.
-  tag: "latest"
-
 imagePullSecrets: []
 nameOverride: ""
 fullnameOverride: ""
@@ -71,8 +65,8 @@ volumeMounts: []
 
 port:
 
-scenario: ""
-
 warnet: ""
 
+archive: ""
+
 args: ""

resources/scenarios/commander.py

@@ -21,7 +21,7 @@
 from test_framework.test_node import TestNode
 from test_framework.util import PortSeed, get_rpc_proxy
 
-WARNET_FILE = Path(os.path.dirname(__file__)) / "warnet.json"
+WARNET_FILE = Path("/warnet.json")
 
 try:
     with open(WARNET_FILE) as file:

resources/scenarios/ln_init.py

@@ -1,12 +1,7 @@
 #!/usr/bin/env python3
 
 from time import sleep
-
-# The base class exists inside the commander container
-try:
-    from commander import Commander
-except ImportError:
-    from resources.scenarios.commander import Commander
+from commander import Commander
 
 
 class LNInit(Commander):
@@ -185,5 +180,5 @@ def funded_lnnodes():
         )
 
 
-if __name__ == "__main__":
+def main():
     LNInit().main()

resources/scenarios/miner_std.py

@@ -1,12 +1,7 @@
 #!/usr/bin/env python3
 
 from time import sleep
-
-# The base class exists inside the commander container
-try:
-    from commander import Commander
-except ImportError:
-    from resources.scenarios.commander import Commander
+from commander import Commander
 
 
 class Miner:
@@ -72,5 +67,5 @@ def run_test(self):
                 sleep(self.options.interval)
 
 
-if __name__ == "__main__":
+def main():
     MinerStd().main()

resources/scenarios/reconnaissance.py

@@ -2,12 +2,7 @@
 
 import socket
 
-# The base class exists inside the commander container when deployed,
-# but requires a relative path inside the python source code for other functions.
-try:
-    from commander import Commander
-except ImportError:
-    from resources.scenarios.commander import Commander
+from commander import Commander
 
 # The entire Bitcoin Core test_framework directory is available as a library
 from test_framework.messages import MSG_TX, CInv, hash256, msg_getdata
@@ -85,5 +80,5 @@ def run_test(self):
         self.log.info(f"Got notfound message from {dstaddr}:{dstport}")
 
 
-if __name__ == "__main__":
+def main():
     Reconnaissance().main()

resources/scenarios/signet_miner.py

@@ -11,11 +11,7 @@
 # we use the authproxy from the test framework.
 ###
 
-# The base class exists inside the commander container
-try:
-    from commander import Commander
-except ImportError:
-    from resources.scenarios.commander import Commander
+from commander import Commander
 
 import json
 import logging
@@ -566,5 +562,5 @@ def get_args(parser):
 
     return args
 
-if __name__ == "__main__":
+def main():
     SignetMinerScenario().main()

test/data/scenario_buggy_failure.py → resources/scenarios/test_buggy_failure.py

@@ -20,5 +20,5 @@ def run_test(self):
         raise Exception("Failed execution!")
 
 
-if __name__ == "__main__":
+def main():
     Failure().main()

test/data/scenario_connect_dag.py → resources/scenarios/test_connect_dag.py

@@ -117,5 +117,5 @@ def assert_connection(self, connector, connectee_index, connection_type: Connect
             raise ValueError("ConnectionType must be of type DNS or IP")
 
 
-if __name__ == "__main__":
+def main():
     ConnectDag().main()

test/data/scenario_p2p_interface.py → resources/scenarios/test_p2p_interface.py

@@ -52,5 +52,5 @@ def run_test(self):
         p2p_block_store.wait_until(lambda: p2p_block_store.blocks[best_block] == 1)
 
 
-if __name__ == "__main__":
+def main():
     GetdataTest().main()

resources/scenarios/tx_flood.py

@@ -1,13 +1,9 @@
 #!/usr/bin/env python3
+
 import threading
 from random import choice, randrange
 from time import sleep
-
-# The base class exists inside the commander container
-try:
-    from commander import Commander
-except ImportError:
-    from resources.scenarios.commander import Commander
+from commander import Commander
 
 
 class TXFlood(Commander):
@@ -70,5 +66,5 @@ def run_test(self):
             sleep(30)
 
 
-if __name__ == "__main__":
+def main():
     TXFlood().main()

ruff.toml

@@ -1,5 +1,5 @@
 extend-exclude = [
-    "resources/images/commander/src/test_framework",
+    "resources/scenarios/test_framework",
     "resources/images/exporter/authproxy.py",
     "resources/scenarios/signet_miner.py",
     "src/test_framework/*",

src/warnet/control.py

@@ -1,9 +1,11 @@
 import base64
+import io
 import json
 import os
 import subprocess
 import sys
 import time
+import zipapp
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from pathlib import Path
 
@@ -167,11 +169,13 @@ def run(scenario_file: str, additional_args: tuple[str]):
     Pass `-- --help` to get individual scenario help
     """
     scenario_path = Path(scenario_file).resolve()
+    scenario_dir = scenario_path.parent
     scenario_name = scenario_path.stem
 
-    with open(scenario_path, "rb") as file:
-        scenario_data = base64.b64encode(file.read()).decode()
+    if additional_args and ("--help" in additional_args or "-h" in additional_args):
+        return subprocess.run([sys.executable, scenario_path, "--help"])
 
+    # Collect tank data for warnet.json
     name = f"commander-{scenario_name.replace('_', '')}-{int(time.time())}"
     namespace = get_default_namespace()
     tankpods = get_mission("tank")
@@ -188,9 +192,39 @@ def run(scenario_file: str, additional_args: tuple[str]):
         for tank in tankpods
     ]
 
-    # Encode warnet data
+    # Encode tank data for warnet.json
     warnet_data = base64.b64encode(json.dumps(tanks).encode()).decode()
 
+    # Create in-memory buffer to store python archive instead of writing to disk
+    archive_buffer = io.BytesIO()
+
+    # No need to copy the entire scenarios/ directory into the archive
+    def filter(path):
+        if any(
+            needle in str(path) for needle in [
+                ".pyc",
+                ".csv",
+                ".DS_Store"
+            ]
+        ):
+            return False
+        return any(
+            needle in str(path) for needle in [
+                "commander.py",
+                "test_framework",
+                scenario_name
+            ]
+        )
+
+    # Compile python archive
+    zipapp.create_archive(
+        source=scenario_dir, target=archive_buffer, main=f"{scenario_name}:main", compressed=True, filter=filter
+    )
+
+    # Encode the binary data as Base64
+    archive_buffer.seek(0)
+    archive_data = base64.b64encode(archive_buffer.read()).decode()
+
     try:
         # Construct Helm command
         helm_command = [
@@ -202,16 +236,14 @@ def run(scenario_file: str, additional_args: tuple[str]):
             "--set",
             f"fullnameOverride={name}",
             "--set",
-            f"scenario={scenario_data}",
-            "--set",
             f"warnet={warnet_data}",
+            "--set",
+            f"archive={archive_data}",
         ]
 
         # Add additional arguments
         if additional_args:
             helm_command.extend(["--set", f"args={' '.join(additional_args)}"])
-            if "--help" in additional_args or "-h" in additional_args:
-                return subprocess.run([sys.executable, scenario_path, "--help"])
 
         helm_command.extend([name, COMMANDER_CHART])
 

test/dag_connection_test.py

@@ -26,7 +26,7 @@ def setup_network(self):
 
     def run_connect_dag_scenario(self):
         self.log.info("Running connect_dag scenario")
-        self.warnet("run test/data/scenario_connect_dag.py")
+        self.warnet("run resources/scenarios/test_connect_dag.py")
         self.wait_for_all_scenarios()
 
 

test/scenarios_test.py

@@ -82,7 +82,7 @@ def run_and_check_miner_scenario_from_file(self):
         self.stop_scenario()
 
     def run_and_check_scenario_from_file(self):
-        scenario_file = "test/data/scenario_p2p_interface.py"
+        scenario_file = "resources/scenarios/test_p2p_interface.py"
         self.log.info(f"Running scenario from: {scenario_file}")
         self.warnet(f"run {scenario_file}")
         self.wait_for_predicate(self.check_scenario_clean_exit)
@@ -94,7 +94,7 @@ def check_regtest_recon(self):
         self.wait_for_predicate(self.check_scenario_clean_exit)
 
     def check_active_count(self):
-        scenario_file = "test/data/scenario_buggy_failure.py"
+        scenario_file = "resources/scenarios/test_buggy_failure.py"
         self.log.info(f"Running scenario from: {scenario_file}")
         self.warnet(f"run {scenario_file}")