support more complicated custom networks

Author: pinheadmz

src/warnet/constants.py

@@ -6,7 +6,7 @@
 # Constants used throughout the project
 # Storing as constants for now but we might want a more sophisticated config management
 # at some point.
-SUPPORTED_TAGS = ["27.0", "26.0", "25.1", "24.2", "23.2", "22.2"]
+SUPPORTED_TAGS = ["29.0", "28.1", "27.0", "26.0", "25.1", "24.2", "23.2", "22.2"]
 DEFAULT_TAG = SUPPORTED_TAGS[0]
 WEIGHTED_TAGS = [
     tag for index, tag in enumerate(reversed(SUPPORTED_TAGS)) for _ in range(index + 1)

src/warnet/graph.py

@@ -4,6 +4,10 @@
 import sys
 from pathlib import Path
 
+from rich import print
+from rich.console import Console
+from rich.table import Table
+
 import click
 import inquirer
 import yaml
@@ -24,9 +28,7 @@ def graph():
 
 
 def custom_graph(
-    num_nodes: int,
-    num_connections: int,
-    version: str,
+    tanks: list,
     datadir: Path,
     fork_observer: bool,
     fork_obs_query_interval: int,
@@ -43,30 +45,39 @@ def custom_graph(
     # Generate network.yaml
     nodes = []
     connections = set()
+    total_count = sum(int(entry["count"]) for entry in tanks)
+    index = 0
 
-    for i in range(num_nodes):
-        node = {"name": f"tank-{i:04d}", "addnode": []}
-
-        # Add round-robin connection
-        next_node = (i + 1) % num_nodes
-        node["addnode"].append(f"tank-{next_node:04d}")
-        connections.add((i, next_node))
-
-        # Add random connections
-        available_nodes = list(range(num_nodes))
-        available_nodes.remove(i)
-        if next_node in available_nodes:
-            available_nodes.remove(next_node)
-
-        for _ in range(min(num_connections - 1, len(available_nodes))):
-            random_node = random.choice(available_nodes)
-            # Avoid circular loops of A -> B -> A
-            if (random_node, i) not in connections:
-                node["addnode"].append(f"tank-{random_node:04d}")
-                connections.add((i, random_node))
-                available_nodes.remove(random_node)
-
-        nodes.append(node)
+    for entry in tanks:
+        for i in range(int(entry["count"])):
+            if ":" in entry["version"] and "/" in entry["version"]:
+                repo, tag = entry["version"].split(":")
+                image = {"repository": repo, "tag": tag}
+            else:
+                image = {"tag": entry["version"]}
+            node = {"name": f"tank-{index:04d}", "addnode": [], "image": image}
+
+            # Add round-robin connection
+            next_node = (index + 1) % total_count
+            node["addnode"].append(f"tank-{next_node:04d}")
+            connections.add((index, next_node))
+
+            # Add random connections
+            available_nodes = list(range(total_count))
+            available_nodes.remove(index)
+            if next_node in available_nodes:
+                available_nodes.remove(next_node)
+
+            for _ in range(min(int(entry["connections"]) - 1, len(available_nodes))):
+                random_node = random.choice(available_nodes)
+                # Avoid circular loops of A -> B -> A
+                if (random_node, index) not in connections:
+                    node["addnode"].append(f"tank-{random_node:04d}")
+                    connections.add((index, random_node))
+                    available_nodes.remove(random_node)
+
+            nodes.append(node)
+            index += 1
 
     network_yaml_data = {"nodes": nodes}
     network_yaml_data["fork_observer"] = {
@@ -86,7 +97,6 @@ def custom_graph(
         "image": {
             "repository": DEFAULT_IMAGE_REPO,
             "pullPolicy": "IfNotPresent",
-            "tag": version
         },
         "defaultConfig":
             f"rpcauth={FORK_OBSERVER_RPCAUTH}\n" +
@@ -108,73 +118,104 @@ def custom_graph(
 
 
 def inquirer_create_network(project_path: Path):
-    # Custom network configuration
-    questions = [
+    network_name_prompt = inquirer.prompt([
         inquirer.Text(
             "network_name",
             message=click.style("Enter your network name", fg="blue", bold=True),
             validate=lambda _, x: len(x) > 0,
-        ),
-        inquirer.List(
-            "nodes",
-            message=click.style("How many nodes would you like?", fg="blue", bold=True),
-            choices=["8", "12", "20", "50", "other"],
-            default="12",
-        ),
-        inquirer.List(
-            "connections",
-            message=click.style(
-                "How many connections would you like each node to have?",
-                fg="blue",
-                bold=True,
-            ),
-            choices=["0", "1", "2", "8", "12", "other"],
-            default="8",
-        ),
-        inquirer.List(
-            "version",
-            message=click.style(
-                "Which version would you like nodes to run by default?", fg="blue", bold=True
-            ),
-            choices=SUPPORTED_TAGS,
-            default=DEFAULT_TAG,
-        ),
-    ]
-
-    net_answers = inquirer.prompt(questions)
-    if net_answers is None:
+        )])
+    if not network_name_prompt:
         click.secho("Setup cancelled by user.", fg="yellow")
         return False
 
-    if net_answers["nodes"] == "other":
-        custom_nodes = inquirer.prompt(
-            [
-                inquirer.Text(
-                    "nodes",
-                    message=click.style("Enter the number of nodes", fg="blue", bold=True),
-                    validate=lambda _, x: int(x) > 0,
-                )
-            ]
-        )
-        if custom_nodes is None:
+    tanks = []
+    while True:
+        table = Table(title="Current Network Population", show_header=True, header_style="magenta")
+        table.add_column("Version", style="cyan")
+        table.add_column("Count", style="green")
+        table.add_column("Connections", style="green")
+
+        for entry in tanks:
+            table.add_row(entry["version"], entry["count"], entry["connections"])
+
+        Console().print(table)
+
+        add_more_prompt = inquirer.prompt([
+            inquirer.List(
+                "add_more",
+                message=click.style(f"How many nodes to add? (0 = done)", fg="blue", bold=True),
+                choices=["0", "4", "8", "12", "20", "50", "other"],
+                default="12")])
+        if not add_more_prompt:
             click.secho("Setup cancelled by user.", fg="yellow")
             return False
-        net_answers["nodes"] = custom_nodes["nodes"]
+        if add_more_prompt["add_more"].startswith("0"):
+            break
 
-    if net_answers["connections"] == "other":
-        custom_connections = inquirer.prompt(
-            [
+        if add_more_prompt["add_more"] == "other":
+            how_many_prompt = inquirer.prompt([
                 inquirer.Text(
-                    "connections",
-                    message=click.style("Enter the number of connections", fg="blue", bold=True),
-                    validate=lambda _, x: int(x) >= 0,
-                )
-            ]
-        )
-        if custom_connections is None:
+                    "how_many",
+                    message=click.style("Enter the number of nodes", fg="blue", bold=True),
+                    validate=lambda _, x: int(x) > 0)])
+            if not how_many_prompt:
+                click.secho("Setup cancelled by user.", fg="yellow")
+                return False
+            how_many = how_many_prompt["how_many"]
+        else:
+            how_many = add_more_prompt["add_more"]
+
+        tank_details_prompt = inquirer.prompt([
+            inquirer.List(
+                "version",
+                message=click.style("Which version would you like to add to network?", fg="blue", bold=True),
+                choices=["other"] + SUPPORTED_TAGS,
+                default=DEFAULT_TAG,
+                ),
+            inquirer.List(
+                "connections",
+                message=click.style(
+                    "How many connections would you like each of these nodes to have?",
+                    fg="blue",
+                    bold=True,
+                ),
+                choices=["0", "1", "2", "8", "12", "other"],
+                default="8",
+            )])
+        if not tank_details_prompt:
             click.secho("Setup cancelled by user.", fg="yellow")
             return False
-        net_answers["connections"] = custom_connections["connections"]
+            break
+        if tank_details_prompt["version"] == "other":
+            custom_version_prompt = inquirer.prompt([
+                inquirer.Text(
+                    "version",
+                    message=click.style("Provide dockerhub repository/image:tag", fg="blue", bold=True),
+                    validate=lambda _, x: "/" in x and ":" in x)])
+            if not custom_version_prompt:
+                click.secho("Setup cancelled by user.", fg="yellow")
+                return False
+            tank_details_prompt["version"] = custom_version_prompt["version"]
+
+        if tank_details_prompt["connections"] == "other":
+            how_many_conn_prompt = inquirer.prompt([
+                inquirer.Text(
+                    "how_many_conn",
+                    message=click.style("Enter the number of connections", fg="blue", bold=True),
+                    validate=lambda _, x: int(x) > 0)])
+            if not how_many_conn_prompt:
+                click.secho("Setup cancelled by user.", fg="yellow")
+                return False
+            how_many_conn = how_many_conn_prompt["how_many_conn"]
+        else:
+            how_many_conn = tank_details_prompt["connections"]
+
+        tanks.append({
+            "version": tank_details_prompt["version"],
+            "count": how_many,
+            "connections": how_many_conn
+        })
+
     fork_observer = click.prompt(
         click.style(
             "\nWould you like to enable fork-observer on the network?", fg="blue", bold=True
@@ -202,12 +243,10 @@ def inquirer_create_network(project_path: Path):
         default=False,
     )
     caddy = fork_observer | logging
-    custom_network_path = project_path / "networks" / net_answers["network_name"]
+    custom_network_path = project_path / "networks" / network_name_prompt["network_name"]
     click.secho("\nGenerating custom network...", fg="yellow", bold=True)
     custom_graph(
-        int(net_answers["nodes"]),
-        int(net_answers["connections"]),
-        net_answers["version"],
+        tanks,
         custom_network_path,
         fork_observer,
         fork_observer_query_interval,

test/graph_test.py

@@ -2,10 +2,13 @@
 
 import json
 import os
+import sys
 
 import pexpect
 from test_base import TestBase
 
+from warnet.process import stream_command
+
 NETWORKS_DIR = "networks"
 
 
@@ -36,15 +39,20 @@ def directory_not_exist(self):
     def directory_exists(self):
         try:
             self.log.info("testing warnet create, dir does exist")
-            self.sut = pexpect.spawn("warnet create")
+            self.sut = pexpect.spawn("warnet create", encoding="utf-8")
+            self.sut.logfile = sys.stdout
             self.sut.expect("name", timeout=30)
             self.sut.sendline("ANewNetwork")
             self.sut.expect("many", timeout=30)
             self.sut.sendline("")
-            self.sut.expect("connections", timeout=30)
-            self.sut.sendline("")
             self.sut.expect("version", timeout=30)
             self.sut.sendline("")
+            self.sut.expect("connections", timeout=30)
+            self.sut.sendline("")
+            self.sut.expect("many", timeout=30)
+            # Up arrow three times: [12] -> 8 -> 4 -> 0 (done)
+            self.sut.sendline("\x1b[A" * 3)
+            self.sut.sendline("\r")
             self.sut.expect("enable fork-observer", timeout=30)
             self.sut.sendline("")
             self.sut.expect("seconds", timeout=30)
@@ -65,7 +73,7 @@ def run_created_network(self):
             f.write(s)
 
         self.log.info("deploying new network")
-        self.warnet("deploy networks/ANewNetwork")
+        stream_command("warnet deploy networks/ANewNetwork")
         self.wait_for_all_tanks_status(target="running")
         debugs = json.loads(self.warnet("bitcoin rpc tank-0000 logging"))
         # set in defaultConfig

test/plugin_test.py

@@ -18,7 +18,7 @@ def __init__(self):
             Path(os.path.dirname(__file__)) / "data" / "network_with_plugins"
         )
         self.plugins_dir = Path(os.path.dirname(__file__)).parent / "resources" / "plugins"
-        self.simln_exec = "plugins/simln/plugin.py"
+        self.simln_exec = self.plugins_dir / "simln" / "plugin.py"
 
     def run_test(self):
         try: