From 7a9a3b2ab0347a2927f207bdb74f1c096c5c1ffb Mon Sep 17 00:00:00 2001
From: boqiu <82121246@qq.com>
Date: Fri, 22 Nov 2024 19:04:02 +0800
Subject: [PATCH] Add python test for TCP connection by shard config

---
 tests/network_tcp_shard_test.py | 76 +++++++++++++++++++++++++++++++++
 1 file changed, 76 insertions(+)
 create mode 100644 tests/network_tcp_shard_test.py

diff --git a/tests/network_tcp_shard_test.py b/tests/network_tcp_shard_test.py
new file mode 100644
index 0000000..8c3896b
--- /dev/null
+++ b/tests/network_tcp_shard_test.py
@@ -0,0 +1,76 @@
+#!/usr/bin/env python3
+
+import os
+import time
+
+from config.node_config import ZGS_KEY_FILE, ZGS_NODEID
+from test_framework.test_framework import TestFramework
+from utility.utils import p2p_port
+
+class NetworkTcpShardTest(TestFramework):
+    """
+    This is to test TCP connection for shard config mismatched peers of UDP discovery.
+    """
+
+    def setup_params(self):
+        # 1 bootnode and 2 community nodes
+        self.num_nodes = 3
+
+        # setup for node 0 as bootnode
+        self.zgs_node_key_files = [ZGS_KEY_FILE]
+        bootnode_port = p2p_port(0)
+        self.zgs_node_configs[0] = {
+            # enable UDP discovery relevant configs
+            "network_enr_address": "127.0.0.1",
+            "network_enr_tcp_port": bootnode_port,
+            "network_enr_udp_port": bootnode_port,
+
+            # disable trusted nodes
+            "network_libp2p_nodes": [],
+
+            # custom shard config
+            "shard_position": "0/4"
+        }
+
+        # setup node 1 & 2 as community nodes
+        bootnodes = [f"/ip4/127.0.0.1/udp/{bootnode_port}/p2p/{ZGS_NODEID}"]
+        for i in range(1, self.num_nodes):
+            self.zgs_node_configs[i] = {
+                # enable UDP discovery relevant configs
+                "network_enr_address": "127.0.0.1",
+                "network_enr_tcp_port": p2p_port(i),
+                "network_enr_udp_port": p2p_port(i),
+
+                # disable trusted nodes and enable bootnodes
+                "network_libp2p_nodes": [],
+                "network_boot_nodes": bootnodes,
+
+                # custom shard config
+                "shard_position": f"{i}/4"
+            }
+
+    def run_test(self):
+        timeout_secs = 10
+
+        for iter in range(timeout_secs):
+            time.sleep(1)
+            self.log.info("==================================== iter %s", iter)
+
+            for i in range(self.num_nodes):
+                info = self.nodes[i].rpc.admin_getNetworkInfo()
+                self.log.info(
+                    "Node[%s] peers: total = %s, banned = %s, disconnected = %s, connected = %s (in = %s, out = %s)",
+                    i, info["totalPeers"], info["bannedPeers"], info["disconnectedPeers"], info["connectedPeers"], info["connectedIncomingPeers"], info["connectedOutgoingPeers"],
+                )
+
+                if i == timeout_secs - 1:
+                    assert info["totalPeers"] == self.num_nodes - 1
+                    assert info["bannedPeers"] == 0
+                    assert info["disconnectedPeers"] == self.num_nodes - 1
+                    assert info["connectedPeers"] == 0
+
+        self.log.info("====================================")
+        self.log.info("All nodes discovered but not connected for each other")
+
+if __name__ == "__main__":
+    NetworkTcpShardTest().main()