use argparse in tests

Author: willcl-ark

.github/workflows/test.yml

@@ -27,7 +27,7 @@ jobs:
           just start
           kubectl wait --for=condition=Ready --timeout=2m pod rpc-0
           kubectl port-forward svc/rpc 9276:9276 &
-      - run: ./test/scenarios_test.py ${{matrix.backend}}
+      - run: ./test/scenarios_test.py --backend=${{matrix.backend}}
   rpc:
     runs-on: ubuntu-latest
     strategy:
@@ -48,7 +48,7 @@ jobs:
           just start
           kubectl wait --for=condition=Ready --timeout=2m pod rpc-0
           kubectl port-forward svc/rpc 9276:9276 &
-      - run: ./test/rpc_test.py ${{matrix.backend}}
+      - run: ./test/rpc_test.py --backend=${{matrix.backend}}
   graph:
     runs-on: ubuntu-latest
     strategy:
@@ -69,7 +69,7 @@ jobs:
           just start
           kubectl wait --for=condition=Ready --timeout=2m pod rpc-0
           kubectl port-forward svc/rpc 9276:9276 &
-      - run: ./test/graph_test.py ${{matrix.backend}}
+      - run: ./test/graph_test.py --backend=${{matrix.backend}}
   ln:
     runs-on: ubuntu-latest
     strategy:
@@ -90,7 +90,7 @@ jobs:
           just start
           kubectl wait --for=condition=Ready --timeout=2m pod rpc-0
           kubectl port-forward svc/rpc 9276:9276 &
-      - run: ./test/ln_test.py ${{matrix.backend}}
+      - run: ./test/ln_test.py --backend=${{matrix.backend}}
   build:
     needs: [scenarios, rpc, graph, ln] # Only run if the others complete successfully as it's slow
     runs-on: ubuntu-latest
@@ -112,4 +112,4 @@ jobs:
           just start
           kubectl wait --for=condition=Ready --timeout=2m pod rpc-0
           kubectl port-forward svc/rpc 9276:9276 &
-      - run: ./test/build_branch_test.py ${{matrix.backend}}
+      - run: ./test/build_branch_test.py --backend=${{matrix.backend}}

src/warnet/server.py

@@ -169,7 +169,7 @@ def setup_rpc(self):
         # Logs
         self.jsonrpc.register(self.logs_grep)
 
-    def get_warnet(self, network: str) -> Warnet:
+    def get_warnet(self, network: str = "warnet") -> Warnet:
         """
         Will get a warnet from the cache if it exists.
         Otherwise it will create the network using from_network() and save it
@@ -268,7 +268,7 @@ def tank_messages(self, network: str, node_a: int, node_b: int) -> str:
             self.logger.error(msg)
             raise ServerError(message=msg) from e
 
-    def network_export(self, network: str) -> str:
+    def network_export(self, network: str = "warnet") -> str:
         """
         Export all data for sim-ln to subdirectory
         """
@@ -283,7 +283,7 @@ def network_export(self, network: str) -> str:
             self.logger.error(msg)
             raise ServerError(message=msg) from e
 
-    def scenarios_available(self, network: str) -> list[tuple]:
+    def scenarios_available(self, network: str = "warnet") -> list[tuple]:
         """
         List available scenarios in the Warnet Test Framework
         """
@@ -415,7 +415,7 @@ def scenarios_stop(self, pid: int, network: str = "warnet") -> str:
             self.logger.error(msg)
             raise ServerError(message=msg)
 
-    def scenarios_list_running(self, network: str) -> list[dict]:
+    def scenarios_list_running(self, network: str = "warnet") -> list[dict]:
         running = [
             {
                 "pid": sc["pid"],

test/rpc_test.py

@@ -1,12 +1,12 @@
 #!/usr/bin/env python3
-
 import os
 from pathlib import Path
 
 from test_base import TestBase
 
 graph_file_path = Path(os.path.dirname(__file__)) / "data" / "v25_x_12.graphml"
 
+
 base = TestBase()
 base.start_server()
 print(base.warcli(f"network start {graph_file_path}"))

test/test_base.py

@@ -1,3 +1,4 @@
+import argparse
 import atexit
 import os
 import sys
@@ -11,6 +12,11 @@
 from warnet.utils import exponential_backoff
 from warnet.warnet import Warnet
 
+parser = argparse.ArgumentParser(description="TestBase args")
+parser.add_argument("--backend", type=str, default="compose", help="Use compose or k8s backend")
+parser.add_argument("--rust_cli", action="store_true", help="Use experimental Rust CLI")
+args = parser.parse_args()
+
 
 class TestBase:
     def __init__(self):
@@ -31,15 +37,11 @@ def __init__(self):
         self.server_thread = None
         self.stop_threads = threading.Event()
         self.network = True
+        self.rust_cli = args.rust_cli
 
         atexit.register(self.cleanup)
 
-        # Default backend
-        self.backend = "compose"
-        # CLI arg overrides env
-        if len(sys.argv) > 1:
-            self.backend = sys.argv[1]
-
+        self.backend = args.backend
         if self.backend not in ["compose", "k8s"]:
             print(f"Invalid backend {self.backend}")
             sys.exit(1)
@@ -78,9 +80,13 @@ def cleanup(self, signum=None, frame=None):
 
     # Execute a warcli RPC using command line (always returns string)
     def warcli(self, str, network=True):
-        cmd = ["warcli"] + str.split()
-        if network:
-            cmd += ["--network", self.network_name]
+        if not self.rust_cli:
+            print("Not using rust cli")
+            cmd = ["warcli"] + str.split()
+            if network:
+                cmd += ["--network", self.network_name]
+        else:
+            cmd = ["target/debug/warcli"] + str.split()
         proc = run(cmd, capture_output=True)
 
         if proc.stderr: