From d6874ff8f33f1f12c5d4eeb88b3a36e47250b260 Mon Sep 17 00:00:00 2001
From: Roman <zajic@zajic.net>
Date: Tue, 4 Jun 2024 23:04:08 +0200
Subject: [PATCH] test: get peers from peer store

---
 src/libs/common.py                            |  4 +++
 src/node/waku_node.py                         |  6 +++++
 .../test_peer_store.py                        | 25 +++++++++++++++++++
 3 files changed, 35 insertions(+)
 create mode 100644 tests/peer_connection_management/test_peer_store.py

diff --git a/src/libs/common.py b/src/libs/common.py
index b7c57cb0..fdc638ef 100644
--- a/src/libs/common.py
+++ b/src/libs/common.py
@@ -38,3 +38,7 @@ def delay(num_seconds):
 
 def gen_step_id():
     return f"{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}__{str(uuid.uuid4())}"
+
+
+def parse_id(peer_info):
+    return peer_info["multiaddr"].split("/")[-1]
diff --git a/src/node/waku_node.py b/src/node/waku_node.py
index 3402a7e8..9a945ffe 100644
--- a/src/node/waku_node.py
+++ b/src/node/waku_node.py
@@ -233,6 +233,12 @@ class WakuNode:
             check_healthy()
         check_ready()
 
+    def get_id(self):
+        try:
+            return self.info_response["listenAddresses"][0].split("/")[-1]
+        except Exception as ex:
+            raise AttributeError(f"Could not find ID in the info call because of error: {str(ex)}")
+
     def get_enr_uri(self):
         try:
             return self.info_response["enrUri"]
diff --git a/tests/peer_connection_management/test_peer_store.py b/tests/peer_connection_management/test_peer_store.py
new file mode 100644
index 00000000..047b5f7d
--- /dev/null
+++ b/tests/peer_connection_management/test_peer_store.py
@@ -0,0 +1,25 @@
+import pytest
+
+from src.libs.common import parse_id
+from src.libs.custom_logger import get_custom_logger
+from src.steps.relay import StepsRelay
+
+logger = get_custom_logger(__name__)
+
+
+@pytest.mark.usefixtures("setup_main_relay_nodes", "setup_optional_relay_nodes")
+class TestPeerStore(StepsRelay):
+    def test_get_peers(self):
+        nodes = [self.node1, self.node2]
+        nodes.extend(self.optional_nodes)
+        ids = []
+        for node in nodes:
+            node_id = node.get_id()
+            ids.append(node_id)
+
+        for i in range(5):
+            others = []
+            for peer_info in nodes[i].get_peers():
+                others.append(parse_id(peer_info))
+
+            assert (i == 0 and len(others) == 4) or (i > 0 and len(others) == 1), f"Some nodes missing in the peer store of node {ids[i]}"