From ca82e26f0aecfb453ff6717a580c81010f88893a Mon Sep 17 00:00:00 2001
From: Matthew Zipkin <pinheadmz@gmail.com>
Date: Mon, 19 Aug 2024 13:30:03 +0200
Subject: [PATCH] add tank option "restart" to graph to disable auto restart

---
 resources/graphs/default.graphml              |  2 +-
 src/warnet/backend/kubernetes_backend.py      |  4 +--
 src/warnet/graph_schema.json                  |  5 +++
 src/warnet/tank.py                            |  1 +
 .../{services.graphml => restart.graphml}     | 18 +++++++---
 test/graph_test.py                            | 33 +++++++++++++------
 6 files changed, 45 insertions(+), 18 deletions(-)
 rename test/data/{services.graphml => restart.graphml} (74%)

diff --git a/resources/graphs/default.graphml b/resources/graphs/default.graphml
index ce84579df..cc9fb1838 100644
--- a/resources/graphs/default.graphml
+++ b/resources/graphs/default.graphml
@@ -18,7 +18,7 @@
   <graph edgedefault="directed">
     <node id="0">
         <data key="version">27.0</data>
-        <data key="bitcoin_config">-uacomment=w0</data>
+        <data key="bitcoin_config">-uacomment=w0 --badconfig</data>
         <data key="exporter">true</data>
         <data key="collect_logs">true</data>
     </node>
diff --git a/src/warnet/backend/kubernetes_backend.py b/src/warnet/backend/kubernetes_backend.py
index f8821888d..b09332f56 100644
--- a/src/warnet/backend/kubernetes_backend.py
+++ b/src/warnet/backend/kubernetes_backend.py
@@ -559,9 +559,7 @@ def create_pod_object(
                 },
             ),
             spec=client.V1PodSpec(
-                # Might need some more thinking on the pod restart policy, setting to Never for now
-                # This means if a node has a problem it dies
-                restart_policy="OnFailure",
+                restart_policy="OnFailure" if tank.restart else "Never",
                 containers=containers,
                 volumes=volumes,
             ),
diff --git a/src/warnet/graph_schema.json b/src/warnet/graph_schema.json
index ac1f7aa9f..a8f768971 100644
--- a/src/warnet/graph_schema.json
+++ b/src/warnet/graph_schema.json
@@ -22,6 +22,11 @@
       "image": {
         "type": "string",
         "comment": "Bitcoin Core Warnet tank image on Dockerhub with the format repository/image:tag"},
+      "restart": {
+        "type": "boolean",
+        "default": true,
+        "comment": "Automatically restart the tank on failure"
+      },
       "bitcoin_config": {
         "type": "string",
         "default": "",
diff --git a/src/warnet/tank.py b/src/warnet/tank.py
index 418dceb6c..57ba0373e 100644
--- a/src/warnet/tank.py
+++ b/src/warnet/tank.py
@@ -46,6 +46,7 @@ def __init__(self, index: int, warnet):
         self.bitcoin_network = warnet.bitcoin_network
         self.version: str = ""
         self.image: str = ""
+        self.restart = True
         self.bitcoin_config = ""
         self.netem = None
         self.exporter = False
diff --git a/test/data/services.graphml b/test/data/restart.graphml
similarity index 74%
rename from test/data/services.graphml
rename to test/data/restart.graphml
index c9e0a0d01..78e959750 100644
--- a/test/data/services.graphml
+++ b/test/data/restart.graphml
@@ -2,6 +2,7 @@
   <key id="services"        attr.name="services"         attr.type="string"   for="graph" />
   <key id="version"         attr.name="version"          attr.type="string"   for="node" />
   <key id="image"           attr.name="image"            attr.type="string"   for="node" />
+  <key id="restart"         attr.name="restart"          attr.type="boolean"  for="node" />
   <key id="bitcoin_config"  attr.name="bitcoin_config"   attr.type="string"   for="node" />
   <key id="tc_netem"        attr.name="tc_netem"         attr.type="string"   for="node" />
   <key id="exporter"        attr.name="exporter"         attr.type="boolean"  for="node" />
@@ -16,12 +17,21 @@
   <key id="source_policy"   attr.name="source_policy"    attr.type="string"   for="edge" />
   <key id="target_policy"   attr.name="target_policy"    attr.type="string"   for="edge" />
   <graph edgedefault="directed">
-    <!-- <data key="services">cadvisor forkobserver addrmanobserver grafana nodeexporter prometheus</data> -->
     <node id="0">
         <data key="version">27.0</data>
-        <data key="bitcoin_config">-uacomment=w0 -debug=validation</data>
-        <data key="exporter">true</data>
-        <data key="ln">lnd</data>
+        <data key="restart">true</data>
+        <data key="bitcoin_config">--badconfigoption</data>
     </node>
+    <node id="1">
+        <data key="version">27.0</data>
+        <data key="restart">false</data>
+        <data key="bitcoin_config">--badconfigoption</data>
+    </node>
+    <node id="2">
+        <data key="version">27.0</data>
+    </node>
+    <edge id="1" source="0" target="1"></edge>
+    <edge id="2" source="1" target="2"></edge>
+    <edge id="3" source="2" target="0"></edge>
   </graph>
 </graphml>
diff --git a/test/graph_test.py b/test/graph_test.py
index 68485f93d..68b52045c 100755
--- a/test/graph_test.py
+++ b/test/graph_test.py
@@ -14,7 +14,6 @@
 class GraphTest(TestBase):
     def __init__(self):
         super().__init__()
-        self.graph_file_path = Path(os.path.dirname(__file__)) / "data" / "services.graphml"
         self.json_file_path = Path(os.path.dirname(__file__)) / "data" / "LN_10.json"
         self.NUM_IMPORTED_NODES = 10
         self.test_dir = tempfile.TemporaryDirectory()
@@ -27,7 +26,7 @@ def run_test(self):
 
         self.start_server()
         try:
-            self.test_graph_with_optional_services()
+            self.test_graph_with_optional_restart()
             self.test_created_graph()
             self.test_imported_graph()
         finally:
@@ -55,16 +54,30 @@ def validate_graph_schema(self):
         self.log.info("Validating graph schema")
         assert "invalid" not in self.warcli(f"graph validate {Path(self.tf_create)}", False)
         assert "invalid" not in self.warcli(f"graph validate {Path(self.tf_import)}", False)
-        assert "invalid" not in self.warcli(f"graph validate {self.graph_file_path}", False)
 
-    def test_graph_with_optional_services(self):
-        self.log.info("Testing graph with optional services...")
-        self.log.info(self.warcli(f"network start {self.graph_file_path}"))
-        self.wait_for_all_tanks_status(target="running")
-        self.wait_for_all_edges()
-        self.warcli("bitcoin rpc 0 getblockcount")
+    def test_graph_with_optional_restart(self):
+        self.log.info("Testing graph with optional restart...")
+        graph_file_path = Path(os.path.dirname(__file__)) / "data" / "restart.graphml"
+        self.log.info(self.warcli(f"network start {graph_file_path}"))
+
+        def expected_statuses():
+            statuses = self.rpc(method="network_status", params={"network": self.network_name})
+            assert len(statuses) == 3
+            # Tanks 0 and 1 have invalid config options that will crash bitcoind
+            for tank in statuses:
+                # auto-restarting forever
+                if tank["tank_index"] == 0 and tank["bitcoin_status"] != "pending":
+                    return False
+                # crash once with error
+                if tank["tank_index"] == 1 and tank["bitcoin_status"] != "failed":
+                    return False
+                # fine
+                if tank["tank_index"] == 2 and tank["bitcoin_status"] != "running":
+                    return False
+            return True
+
+        self.wait_for_predicate(expected_statuses)
 
-        self.log.info("Checking services...")
         self.warcli("network down")
         self.wait_for_all_tanks_status(target="stopped")