From b11826ebd9b3e3b0eb827aa59486fa9b2c590340 Mon Sep 17 00:00:00 2001
From: Roman <zajic@zajic.net>
Date: Tue, 30 Apr 2024 18:33:17 +0800
Subject: [PATCH] fix: add timeout param to ensure_ready

---
 src/node/waku_node.py   | 17 ++++++++++-------
 tests/relay/test_rln.py |  4 ++--
 2 files changed, 12 insertions(+), 9 deletions(-)

diff --git a/src/node/waku_node.py b/src/node/waku_node.py
index 3effa16e..62d0595b 100644
--- a/src/node/waku_node.py
+++ b/src/node/waku_node.py
@@ -48,7 +48,7 @@ class WakuNode:
         logger.debug(f"WakuNode instance initialized with log path {self._log_path}")
 
     @retry(stop=stop_after_delay(5), wait=wait_fixed(0.1), reraise=True)
-    def start(self, **kwargs):
+    def start(self, wait_for_node_sec=10, **kwargs):
         logger.debug("Starting Node...")
         self._docker_manager.create_network()
         self._ext_ip = self._docker_manager.generate_random_ext_ip()
@@ -123,7 +123,7 @@ class WakuNode:
         DS.waku_nodes.append(self)
         delay(1)  # if we fire requests to soon after starting the node will sometimes fail to start correctly
         try:
-            self.ensure_ready()
+            self.ensure_ready(timeout_duration=wait_for_node_sec)
         except Exception as ex:
             logger.error(f"REST service did not become ready in time: {ex}")
             raise
@@ -186,12 +186,15 @@ class WakuNode:
             logger.debug(f"Unpause container with id {self._container.short_id}")
             self._container.unpause()
 
-    @retry(stop=stop_after_delay(600), wait=wait_fixed(1), reraise=True)
-    def ensure_ready(self):
-        self.info_response = self.info()
-        logger.info("REST service is ready !!")
+    def ensure_ready(self, timeout_duration=10):
+        @retry(stop=stop_after_delay(timeout_duration), wait=wait_fixed(0.1), reraise=True)
+        def check_ready(node=self):
+            node.info_response = node.info()
+            logger.info("REST service is ready !!")
 
-    @retry(stop=stop_after_delay(600), wait=wait_fixed(1), reraise=True)
+        check_ready()
+
+    @retry(stop=stop_after_delay(10), wait=wait_fixed(1), reraise=True)
     def ensure_healthy(self):
         self.health_response = self.health()
         logger.info("Node is healthy !!")
diff --git a/tests/relay/test_rln.py b/tests/relay/test_rln.py
index 48a94208..e0c4e748 100644
--- a/tests/relay/test_rln.py
+++ b/tests/relay/test_rln.py
@@ -103,7 +103,7 @@ class TestRelayRLN(StepsRLN, StepsRelay):
     @pytest.mark.skip(reason="exceeding timeout ATM, need to check with devs")
     @pytest.mark.timeout(600)
     def test_valid_payloads_dynamic_at_slow_rate(self):
-        self.setup_main_rln_relay_nodes(rln_relay_dynamic="true")
+        self.setup_main_rln_relay_nodes(rln_relay_dynamic="true", wait_for_node_sec=600)
         self.subscribe_main_relay_nodes()
         failed_payloads = []
         for payload in SAMPLE_INPUTS:
@@ -120,7 +120,7 @@ class TestRelayRLN(StepsRLN, StepsRelay):
     @pytest.mark.skip(reason="exceeding timeout ATM, need to check with devs")
     @pytest.mark.timeout(600)
     def test_valid_payloads_dynamic_at_spam_rate(self):
-        self.setup_main_rln_relay_nodes(rln_relay_dynamic="true")
+        self.setup_main_rln_relay_nodes(rln_relay_dynamic="true", wait_for_node_sec=600)
         self.subscribe_main_relay_nodes()
         previous = math.trunc(time())
         for i, payload in enumerate(SAMPLE_INPUTS):