|
1 | 1 | import logging |
2 | | -import time |
3 | | -import pytest |
4 | 2 | import threading |
5 | 3 |
|
6 | | -from clients.status_backend import StatusBackend |
7 | | - |
8 | | -DISCOVERY_TIMEOUT_SEC = 30 |
9 | | - |
10 | | - |
11 | | -def _all_nodes_discovered(nodes: dict[str, StatusBackend], known_peers: dict[str, str]) -> bool: |
12 | | - """Check if all nodes discovered each other""" |
13 | | - for peer_id, node in nodes.items(): |
14 | | - if node is None: |
15 | | - return False |
16 | | - peers = node.wakuext_service.peers() |
17 | | - |
18 | | - # Use shorter names for logging |
19 | | - peer_names = [known_peers.get(peer, peer[-5:]) for peer in peers] |
20 | | - logging.info(f"Node {known_peers[peer_id]} peers: {peer_names}") |
| 4 | +import pytest |
| 5 | +from tenacity import retry, stop_after_delay, wait_fixed |
21 | 6 |
|
22 | | - for known_node in known_peers.keys(): |
23 | | - if known_node == peer_id: |
24 | | - continue |
25 | | - if known_node not in peers: |
26 | | - return False |
27 | | - return True |
| 7 | +from clients.status_backend import StatusBackend |
28 | 8 |
|
29 | 9 |
|
30 | 10 | @pytest.mark.rpc |
31 | 11 | class TestDiscovery: |
32 | | - |
33 | 12 | def test_discovery(self, backend_new_profile): |
34 | | - nodes_count = 3 |
| 13 | + # Desired topology: 3 full nodes, 2 light nodes |
| 14 | + full_count = 3 |
| 15 | + light_count = 2 |
| 16 | + total_nodes = full_count + light_count |
| 17 | + |
35 | 18 | nodes: dict[str, StatusBackend] = {} |
| 19 | + full_ids: set[str] = set() |
36 | 20 |
|
37 | 21 | known_nodes = { |
| 22 | + # Boot nodes available in the fleet |
38 | 23 | "16Uiu2HAm3vFYHkGRURyJ6F7bwDyzMLtPEuCg4DU89T7km2u8Fjyb": "boot-1", |
39 | 24 | "16Uiu2HAmCDqxtfF1DwBqs7UJ4TgSnjoh6j1RtE1hhQxLLao84jLi": "store", |
40 | 25 | } |
41 | 26 |
|
42 | | - def create_node(node_index: int): |
43 | | - """Function to run in each thread - waits for wakuv2.peerstats signal""" |
44 | | - backend = backend_new_profile(f"node_{node_index}") |
| 27 | + def create_node(node_index: int, waku_light_client: bool): |
| 28 | + backend = backend_new_profile( |
| 29 | + f"node_{node_index}", |
| 30 | + # We mix full and light clients in a single test. |
| 31 | + # Light clients (edge nodes) will not discover each, |
| 32 | + # but they must discover full clients. |
| 33 | + waku_light_client=waku_light_client, |
| 34 | + # Force the container to only have the docker compose network. |
| 35 | + # Otherwise, advertised ENRs may contain the wrong IP address. |
| 36 | + bridge_network=False, |
| 37 | + ) |
45 | 38 | peer_id = backend.wakuext_service.peer_id() |
46 | 39 | known_nodes[peer_id] = f"backend_{node_index}" |
| 40 | + if not waku_light_client: |
| 41 | + full_ids.add(peer_id) |
47 | 42 | nodes[peer_id] = backend |
48 | | - logging.info(f"Backend {node_index} ready. Peer ID: {peer_id}") |
49 | | - |
50 | | - # Run threads, each waiting for wakuv2.peerstats signal |
51 | | - logging.info("Starting threads to wait for wakuv2.peerstats signals...") |
| 43 | + info = f"Peer ID: {peer_id}" |
| 44 | + info += f", URL: {backend.url}" |
| 45 | + if backend.container: |
| 46 | + info += f", Container: {backend.container.short_id()}" |
| 47 | + info += f", Type: {'full' if not waku_light_client else 'light'}" |
| 48 | + logging.info(f"Backend {node_index} ready. {info}") |
| 49 | + |
| 50 | + # Start threads to create nodes |
| 51 | + logging.info("Starting threads to create full and light clients...") |
52 | 52 | threads = [] |
53 | 53 |
|
54 | | - for i in range(nodes_count): |
55 | | - thread = threading.Thread(target=create_node, args=(i,)) |
56 | | - thread.daemon = True |
57 | | - thread.start() |
58 | | - threads.append(thread) |
59 | | - |
60 | | - for thread in threads: |
61 | | - thread.join() |
62 | | - |
63 | | - assert len(nodes.keys()) == nodes_count, "Not all nodes created" |
64 | | - |
65 | | - # Wait for all nodes to discover each other |
66 | | - start_time = time.time() |
67 | | - while time.time() - start_time < DISCOVERY_TIMEOUT_SEC: |
68 | | - if _all_nodes_discovered(nodes, known_nodes): |
69 | | - return |
70 | | - time.sleep(0.5) |
71 | | - |
72 | | - assert False, f"Nodes failed to discover each other within {DISCOVERY_TIMEOUT_SEC} seconds" |
| 54 | + # First start full nodes |
| 55 | + for i in range(full_count): |
| 56 | + t = threading.Thread(target=create_node, args=(i, False)) |
| 57 | + t.daemon = True |
| 58 | + t.start() |
| 59 | + threads.append(t) |
| 60 | + |
| 61 | + # Then start light nodes |
| 62 | + for i in range(full_count, total_nodes): |
| 63 | + t = threading.Thread(target=create_node, args=(i, True)) |
| 64 | + t.daemon = True |
| 65 | + t.start() |
| 66 | + threads.append(t) |
| 67 | + |
| 68 | + for t in threads: |
| 69 | + t.join() |
| 70 | + |
| 71 | + assert len(nodes.keys()) == total_nodes, "Not all nodes created" |
| 72 | + |
| 73 | + # Build sets for expectations |
| 74 | + all_node_ids = set(nodes.keys()) |
| 75 | + boot_ids = {k for k, v in known_nodes.items() if v in ("boot-1", "store")} |
| 76 | + |
| 77 | + # Pre-build expected peers per node (static expectations) |
| 78 | + expected_by_node: dict[str, set[str]] = {} |
| 79 | + for pid in nodes.keys(): |
| 80 | + if pid in full_ids: |
| 81 | + # Full clients should discover all other clients and boot nodes |
| 82 | + expected = (all_node_ids | boot_ids) - {pid} |
| 83 | + else: |
| 84 | + # Light clients should discover only full clients and boot nodes |
| 85 | + expected = boot_ids | full_ids |
| 86 | + expected_by_node[pid] = expected |
| 87 | + |
| 88 | + # Check using tenacity.retry, logs peers on each iteration |
| 89 | + @retry(stop=stop_after_delay(120), wait=wait_fixed(1), reraise=True) |
| 90 | + def assert_discovery(): |
| 91 | + all_ok = True |
| 92 | + |
| 93 | + for peer_id, node in nodes.items(): |
| 94 | + # Fetch peers once per node per attempt |
| 95 | + peers = set(node.wakuext_service.peers()) |
| 96 | + |
| 97 | + # Log peer names |
| 98 | + peer_names = [known_nodes.get(peer, peer[-5:]) for peer in peers] |
| 99 | + logging.info(f"Checking node {known_nodes[peer_id]} peers: {peer_names}") |
| 100 | + |
| 101 | + if peers >= expected_by_node[peer_id]: |
| 102 | + logging.info(f"Node {known_nodes[peer_id]} discovered peers as expected") |
| 103 | + nodes.pop(peer_id) |
| 104 | + continue |
| 105 | + |
| 106 | + all_ok = False |
| 107 | + |
| 108 | + assert all_ok, "Not all nodes discovered as expected" |
| 109 | + |
| 110 | + assert_discovery() |
0 commit comments