From df2e5e32b2839130de5fb287d9967a8acbc3b8d4 Mon Sep 17 00:00:00 2001
From: Roman <zajic@zajic.net>
Date: Wed, 5 Jun 2024 21:09:01 +0200
Subject: [PATCH] fix: add go_waku support

---
 src/libs/common.py                            | 11 +++---
 .../test_peer_store.py                        | 34 +++++++++----------
 2 files changed, 23 insertions(+), 22 deletions(-)

diff --git a/src/libs/common.py b/src/libs/common.py
index f8b31f8a..fcad2b53 100644
--- a/src/libs/common.py
+++ b/src/libs/common.py
@@ -40,12 +40,15 @@ def gen_step_id():
     return f"{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}__{str(uuid.uuid4())}"
 
 
-def peer_info2id(peer):
-    return peer_info2multiaddr(peer).split("/")[-1]
+def peer_info2id(peer, is_nwaku=True):
+    return peer_info2multiaddr(peer, is_nwaku).split("/")[-1]
 
 
-def peer_info2multiaddr(peer):
-    return peer["multiaddr"]
+def peer_info2multiaddr(peer, is_nwaku=True):
+    if is_nwaku:
+        return peer["multiaddr"]
+    else:
+        return peer["multiaddrs"][0]
 
 
 def multiaddr2id(multiaddr):
diff --git a/tests/peer_connection_management/test_peer_store.py b/tests/peer_connection_management/test_peer_store.py
index 43fe5526..638c0cb0 100644
--- a/tests/peer_connection_management/test_peer_store.py
+++ b/tests/peer_connection_management/test_peer_store.py
@@ -12,6 +12,7 @@ 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()
@@ -20,36 +21,33 @@ class TestPeerStore(StepsRelay):
         for i in range(5):
             others = []
             for peer_info in nodes[i].get_peers():
-                others.append(peer_info2id(peer_info))
+                peer_id = peer_info2id(peer_info, nodes[i].is_nwaku())
+                others.append(peer_id)
 
-            assert (i == 0 and len(others) == 4) or (i > 0 and len(others) == 1), f"Some nodes missing in the peer store of Node ID {ids[i]}"
+            assert (i == 0 and len(others) == 4) or (i > 0 and len(others) >= 1), f"Some nodes missing in the peer store of Node ID {ids[i]}"
 
     def test_add_peers(self):
         nodes = [self.node1, self.node2]
         nodes.extend(self.optional_nodes)
 
-        # Get peers 1-4
-        peers_info = nodes[0].get_peers()
-        assert len(peers_info) == 4, f"Some nodes missing in the peer store of Node 1"
+        peers_multiaddr = set()
+        for i in range(2):
+            for peer_info in nodes[i].get_peers():
+                multiaddr = peer_info2multiaddr(peer_info, nodes[i].is_nwaku())
+                peers_multiaddr.add(multiaddr)
 
-        # Get peer 0
-        peers_info.extend(nodes[1].get_peers())
-        assert len(peers_info) == 5, f"Node 1 missing in the peer store of Node 2"
-
-        # Convert to multi addresses
-        peers_multiaddr = []
-        for peer in peers_info:
-            multiaddr = peer_info2multiaddr(peer)
-            logger.debug(f"Peer info {peer}")
-            logger.debug(f"Peer multi address {multiaddr}")
-            peers_multiaddr.append(multiaddr)
+        assert len(peers_multiaddr) == 5, f"Exactly 5 multi addresses are expected"
 
         # Add peers one by one excluding self for Nodes 2-5
         for i in range(1, 5):
-            for peer in peers_multiaddr:
+            for peer in list(peers_multiaddr):
                 if nodes[i].get_id() != multiaddr2id(peer):
                     try:
-                        nodes[i].add_peers([peer])
+                        if nodes[i].is_nwaku():
+                            nodes[i].add_peers([peer])
+                        else:
+                            peer_info = {"multiaddr": peer, "shards": [0], "protocols": ["/vac/waku/relay/2.0.0"]}
+                            nodes[i].add_peers(peer_info)
                     except Exception as ex:
                         logger.error(f"Failed to add peer to Node {i} peer store: {ex}")
                         raise