add create-network cli command

pyproject.toml

@@ -19,6 +19,7 @@ dependencies = [
     "rich==13.7.1",
     "tabulate==0.9.0",
     "PyYAML==6.0.2",
+    "pexpect==4.9.0",
 ]
 
 [project.scripts]

src/warnet/graph.py

@@ -1,6 +1,13 @@
+import os
+import random
+from importlib.resources import files
 from pathlib import Path
 
 import click
+import inquirer
+import yaml
+
+from .util import DEFAULT_TAG, SUPPORTED_TAGS
 
 
 @click.group(name="graph", hidden=True)
@@ -17,3 +24,159 @@ def import_json(infile: Path, outfile: Path, cb: str, ln_image: str):
     Returns XML file as string with or without --outfile option.
     """
     raise Exception("Not Implemented")
+
+
+def custom_graph(num_nodes: int, num_connections: int, version: str, datadir: Path):
+    datadir.mkdir(parents=False, exist_ok=False)
+    # Generate network.yaml
+    nodes = []
+    connections = set()
+
+    for i in range(num_nodes):
+        node = {"name": f"tank-{i:04d}", "connect": [], "image": {"tag": version}}
+
+        # Add round-robin connection
+        next_node = (i + 1) % num_nodes
+        node["connect"].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))):
+            # trunk-ignore(bandit/B311)
+            random_node = random.choice(available_nodes)
+            # Avoid circular loops of A -> B -> A
+            if (random_node, i) not in connections:
+                node["connect"].append(f"tank-{random_node:04d}")
+                connections.add((i, random_node))
+                available_nodes.remove(random_node)
+
+        nodes.append(node)
+
+    network_yaml_data = {"nodes": nodes}
+
+    with open(os.path.join(datadir, "network.yaml"), "w") as f:
+        yaml.dump(network_yaml_data, f, default_flow_style=False)
+
+    # Generate defaults.yaml
+    default_yaml_path = files("resources.networks").joinpath("6_node_bitcoin/node-defaults.yaml")
+    with open(str(default_yaml_path)) as f:
+        defaults_yaml_content = f.read()
+
+    with open(os.path.join(datadir, "node-defaults.yaml"), "w") as f:
+        f.write(defaults_yaml_content)
+
+    click.echo(
+        f"Project '{datadir}' has been created with 'network.yaml' and 'node-defaults.yaml'."
+    )
+
+
+def inquirer_create_network(project_path: Path):
+    # Custom network configuration
+    questions = [
+        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:
+        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:
+            click.secho("Setup cancelled by user.", fg="yellow")
+            return False
+        net_answers["nodes"] = custom_nodes["nodes"]
+
+    if net_answers["connections"] == "other":
+        custom_connections = 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:
+            click.secho("Setup cancelled by user.", fg="yellow")
+            return False
+        net_answers["connections"] = custom_connections["connections"]
+
+    custom_network_path = project_path / "networks" / net_answers["network_name"]
+    click.secho("\nGenerating custom network...", fg="yellow", bold=True)
+    custom_graph(
+        int(net_answers["nodes"]),
+        int(net_answers["connections"]),
+        net_answers["version"],
+        custom_network_path,
+    )
+    return custom_network_path
+
+
+@click.command()
+def create_network():
+    """Create a new warnet network"""
+    try:
+        project_path = Path(os.getcwd())
+        # Check if the project has a networks directory
+        if not (project_path / "networks").exists():
+            click.secho(
+                "The current directory does not have a 'networks' directory. Please run 'warnet init' or 'warnet create' first.",
+                fg="red",
+                bold=True,
+            )
+            return False
+        custom_network_path = inquirer_create_network(project_path)
+        click.secho("\nNew network created successfully!", fg="green", bold=True)
+        click.echo("\nRun the following command to deploy this network:")
+        click.echo(f"warnet deploy {custom_network_path}")
+    except Exception as e:
+        click.echo(f"{e}\n\n")
+        click.secho(f"An error occurred while creating a new network:\n\n{e}\n\n", fg="red")
+        click.secho(
+            "Please report the above context to https://github.com/bitcoin-dev-project/warnet/issues",
+            fg="yellow",
+        )
+        return False

src/warnet/main.py

@@ -1,6 +1,5 @@
 import json
 import os
-import random
 import subprocess
 import sys
 from importlib.resources import files
@@ -18,11 +17,10 @@
 from .bitcoin import bitcoin
 from .control import down, run, stop
 from .deploy import deploy as deploy_command
-from .graph import graph
+from .graph import create_network, graph, inquirer_create_network
 from .image import image
 from .network import copy_network_defaults, copy_scenario_defaults
 from .status import status as status_command
-from .util import DEFAULT_TAG, SUPPORTED_TAGS
 
 QUICK_START_PATH = files("resources.scripts").joinpath("quick_start.sh")
 
@@ -41,6 +39,7 @@ def cli():
 cli.add_command(stop)
 cli.add_command(down)
 cli.add_command(run)
+cli.add_command(create_network)
 
 
 @cli.command()
@@ -112,89 +111,13 @@ def quickstart():
             click.echo(f"warcli deploy {proj_answers['project_path']}/networks/6_node_bitcoin")
             return True
         answers.update(proj_answers)
-
-        # Custom network configuration
-        questions = [
-            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:
-            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:
-                click.secho("Setup cancelled by user.", fg="yellow")
-                return False
-            net_answers["nodes"] = custom_nodes["nodes"]
-
-        if net_answers["connections"] == "other":
-            custom_connections = 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:
-                click.secho("Setup cancelled by user.", fg="yellow")
-                return False
-            net_answers["connections"] = custom_connections["connections"]
-        answers.update(net_answers)
-
         click.secho("\nCreating project structure...", fg="yellow", bold=True)
         project_path = Path(os.path.expanduser(proj_answers["project_path"]))
         create_warnet_project(project_path)
+
         click.secho("\nGenerating custom network...", fg="yellow", bold=True)
-        custom_network_path = project_path / "networks" / answers["network_name"]
-        custom_graph(
-            int(answers["nodes"]),
-            int(answers["connections"]),
-            answers["version"],
-            custom_network_path,
-        )
+        custom_network_path = inquirer_create_network(project_path)
+
         click.secho("\nSetup completed successfully!", fg="green", bold=True)
         click.echo("\nRun the following command to deploy this network:")
         click.echo(f"warnet deploy {custom_network_path}")
@@ -357,51 +280,3 @@ def logs(pod_name: str, follow: bool):
 
 if __name__ == "__main__":
     cli()
-
-
-def custom_graph(num_nodes: int, num_connections: int, version: str, datadir: Path):
-    datadir.mkdir(parents=False, exist_ok=False)
-    # Generate network.yaml
-    nodes = []
-    connections = set()
-
-    for i in range(num_nodes):
-        node = {"name": f"tank-{i:04d}", "connect": [], "image": {"tag": version}}
-
-        # Add round-robin connection
-        next_node = (i + 1) % num_nodes
-        node["connect"].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["connect"].append(f"tank-{random_node:04d}")
-                connections.add((i, random_node))
-                available_nodes.remove(random_node)
-
-        nodes.append(node)
-
-    network_yaml_data = {"nodes": nodes}
-
-    with open(os.path.join(datadir, "network.yaml"), "w") as f:
-        yaml.dump(network_yaml_data, f, default_flow_style=False)
-
-    # Generate defaults.yaml
-    default_yaml_path = files("resources.networks").joinpath("6_node_bitcoin/node-defaults.yaml")
-    with open(str(default_yaml_path)) as f:
-        defaults_yaml_content = f.read()
-
-    with open(os.path.join(datadir, "node-defaults.yaml"), "w") as f:
-        f.write(defaults_yaml_content)
-
-    click.echo(
-        f"Project '{datadir}' has been created with 'network.yaml' and 'node-defaults.yaml'."
-    )

test/graph_test.py

@@ -0,0 +1,44 @@
+#!/usr/bin/env python3
+
+import os
+import shutil
+
+import pexpect
+from test_base import TestBase
+
+NETWORKS_DIR = "networks"
+
+
+class GraphTest(TestBase):
+    def __init__(self):
+        super().__init__()
+
+    def run_test(self):
+        try:
+            self.directory_not_exist()
+            os.mkdir(NETWORKS_DIR)
+            self.directory_exists()
+
+        finally:
+            shutil.rmtree(NETWORKS_DIR) if os.path.exists(NETWORKS_DIR) else None
+
+    def directory_not_exist(self):
+        self.sut = pexpect.spawn("warnet create-network")
+        self.sut.expect("init", timeout=5)
+
+    def directory_exists(self):
+        self.sut = pexpect.spawn("warnet create-network")
+        self.sut.expect("name", timeout=1)
+        self.sut.sendline("ANewNetwork")
+        self.sut.expect("many", timeout=1)
+        self.sut.sendline("")
+        self.sut.expect("connections", timeout=1)
+        self.sut.sendline("")
+        self.sut.expect("version", timeout=1)
+        self.sut.sendline("")
+        self.sut.expect("successfully", timeout=5)
+
+
+if __name__ == "__main__":
+    test = GraphTest()
+    test.run_test()