From 5e233fe2ed1d449dbf4a97f489274335fb0abb6e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micha=C5=82=20Kudela?= <kudmich@syncad.com>
Date: Tue, 26 Nov 2024 15:55:56 +0100
Subject: [PATCH] Handle `api mode` notification
---
package/test_tools/__private/node.py | 19 ++++++++++++++++---
.../node_notification_handler.py | 3 +++
2 files changed, 19 insertions(+), 3 deletions(-)
diff --git a/package/test_tools/__private/node.py b/package/test_tools/__private/node.py
index b07c43005..d72dbff11 100644
--- a/package/test_tools/__private/node.py
+++ b/package/test_tools/__private/node.py
@@ -375,11 +375,13 @@ class Node(BaseNode, ScopedObject):
self.__produced_files = True
if not exit_before_synchronization and exit_at_block is None:
- self.__wait_for_synchronization(deadline, timeout, wait_for_live)
+ self.__wait_for_synchronization(deadline, timeout, wait_for_live, stop_at_block)
self.__log_run_summary()
- def __wait_for_synchronization(self, deadline: float, timeout: float, wait_for_live: bool) -> None:
+ def __wait_for_synchronization(
+ self, deadline: float, timeout: float, wait_for_live: bool, stop_at_block: int | None
+ ) -> None:
wait_for_event(
self.__notifications.handler.synchronization_started_event,
deadline=deadline,
@@ -398,7 +400,9 @@ class Node(BaseNode, ScopedObject):
exception_message="WS server didn't start listening on time.",
)
- if wait_for_live:
+ if stop_at_block is not None or "queen" in self.config.plugin:
+ self.wait_for_api_mode(timeout=timeout)
+ elif wait_for_live:
self.wait_for_live_mode(timeout=timeout)
wait_for_event(
@@ -615,6 +619,15 @@ class Node(BaseNode, ScopedObject):
exception_message=f"{self.get_name()}: Live mode not activated on time.",
)
+ def wait_for_api_mode(self, timeout: float = math.inf) -> None:
+ assert timeout >= 0
+ deadline = time.time() + timeout
+ wait_for_event(
+ self.__notifications.handler.api_mode_entered_event,
+ deadline=deadline,
+ exception_message=f"{self.get_name()}: API mode not activated on time.",
+ )
+
def get_number_of_forks(self) -> int:
return self.__notifications.handler.number_of_forks
diff --git a/package/test_tools/__private/notifications/node_notification_handler.py b/package/test_tools/__private/notifications/node_notification_handler.py
index 7763b2388..e0bf25fa7 100644
--- a/package/test_tools/__private/notifications/node_notification_handler.py
+++ b/package/test_tools/__private/notifications/node_notification_handler.py
@@ -44,6 +44,7 @@ class NodeNotificationHandler(HivedNotificationHandler):
self.replay_finished_event = Event()
self.snapshot_dumped_event = Event()
+ self.api_mode_entered_event = Event()
self.switch_fork_event = Event()
self.number_of_forks = 0
@@ -60,6 +61,8 @@ class NodeNotificationHandler(HivedNotificationHandler):
self.live_mode_entered_event.set()
case "chain API ready":
self.chain_api_ready_event.set()
+ case "entering API mode":
+ self.api_mode_entered_event.set()
async def on_http_webserver_bind(self, notification: Notification[WebserverListening]) -> None:
self.http_endpoint = HttpUrl(self.__combine_url_string_from_notification(notification), protocol="http")
--
GitLab