diff --git a/examples/kademlia/grid_topology_example.py b/examples/kademlia/grid_topology_example.py new file mode 100644 index 000000000..004707e1f --- /dev/null +++ b/examples/kademlia/grid_topology_example.py @@ -0,0 +1,170 @@ +#!/usr/bin/env python + +""" +Grid Topology DHT Example + +This example demonstrates how to use the Grid Topology routing table +with py-libp2p's Kademlia DHT. Grid topology provides a fixed 256-bucket +structure based on Common Prefix Length (CPL) indexing, matching cpp-libp2p. + +Grid Topology is useful when: +- You need better interoperability with cpp-libp2p or Go-libp2p +- You prefer a simpler, more predictable bucket structure +- You need explicit peer state management (temporary vs permanent peers) +""" + +import logging + +from libp2p.kad_dht.grid_routing_table import GridRoutingTable, NodeId +from libp2p.kad_dht.grid_topology_config import ( + get_default_config, +) +from libp2p.peer.id import ID +from libp2p.peer.peerinfo import PeerInfo + +# Configure logging +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + + +def example_basic_grid_topology(): + """Basic usage of Grid Topology routing table.""" + print("\nBasic Grid Topology Usage\n") + + # Create local node ID + local_id = ID.from_base58("QmaCpDMGvV2BGHeYERUEnRQAwe5CcqarqmtA7xNXT92p2") + print(f"Local Node ID: {local_id}") + + # Create grid routing table with default config + config = get_default_config() + rt = GridRoutingTable(local_id, max_bucket_size=config.max_bucket_size) + print("Created Grid Topology with 256 fixed buckets\n") + + # Create some test peers + peer_ids = [ + ID.from_base58("QmZLaXk3bbiHgVK3zp5A8n2DEuvMZFRv1GAjTrSvZuLnFr"), + ID.from_base58("QmdgB6x6xfBLvV9VwSPj9D7aHCmXhvEVVBn2CUEbNJnTg"), + ID.from_base58("QmYPp4CUQpRcpHBqWgmQ4TjqW9AHJ3qLhQV5d7s3d8hMJx"), + ] + + # Add peers to the routing table + print("Adding peers to routing table:") + for peer_id in peer_ids: + peer_info = PeerInfo(peer_id, []) + success = rt.update(peer_id, peer_info=peer_info, is_permanent=True) + + if success: + node_id = NodeId(peer_id) + bucket_index = rt._get_bucket_index(node_id) + print(f" ✓ Added peer to bucket {bucket_index}") + else: + print(" ✗ Failed to add peer") + + print(f"\nTotal peers in routing table: {rt.size()}\n") + return rt, local_id, peer_ids + + +def example_peer_lookup(rt, peer_ids): + """Demonstrate peer lookup in grid topology.""" + print("Peer Lookup\n") + + for peer_id in peer_ids: + if rt.contains(peer_id): + print(f"✓ Peer {str(peer_id)[:16]}... found in routing table") + + # Get peer info from the bucket + node_id = NodeId(peer_id) + bucket_index = rt._get_bucket_index(node_id) + if bucket_index is not None: + bucket = rt.get_bucket(bucket_index) + if bucket: + peer_info = bucket.get_peer_info(peer_id) + if peer_info: + print(f" Peer info: {peer_info}\n") + else: + print(f"✗ Peer {str(peer_id)[:16]}... NOT found\n") + + +def example_bucket_distribution(rt): + """Show how peers are distributed across buckets.""" + print("Bucket Distribution\n") + + stats = rt.get_bucket_stats() + print(f"Total peers: {stats['total_peers']}") + print(f"Total buckets: {stats['total_buckets']}") + print(f"Non-empty buckets: {stats['non_empty_buckets']}") + + if stats["total_peers"] > 0: + avg = stats["total_peers"] / max(1, stats["non_empty_buckets"]) + print(f"Average peers per bucket: {avg:.1f}\n") + + +def example_grid_topology_advantages(): + """Explain advantages of grid topology.""" + print("Grid Topology Advantages\n") + + print("1. Fixed 256 Buckets") + print(" - One bucket per bit in 256-bit ID space") + print(" - No dynamic splitting complexity") + print(" - Predictable bucket structure\n") + + print("2. CPL-Based Indexing") + print(" - Bucket index = 255 - CPL(local_id, peer_id)") + print(" - CPL = Common Prefix Length") + print(" - Simpler than XOR range-based indexing\n") + + print("3. Explicit Peer States") + print(" - Permanent peers: Critical connections") + print(" - Temporary peers: Can be replaced") + print(" - Better control over peer replacement\n") + + print("4. Cross-Implementation Compatibility") + print(" - Matches cpp-libp2p specification") + print(" - Compatible with Go-libp2p grid topology") + print(" - Better interoperability\n") + + +def example_compare_with_standard_kbucket(): + """Show how grid topology differs from standard k-bucket.""" + print("Grid Topology vs Standard K-Bucket\n") + + print("Standard K-Bucket:") + print(" - Dynamic buckets (starts with 1, splits as needed)") + print(" - XOR range-based indexing") + print(" - Implicit LRU peer ordering") + print(" - Good for general use cases\n") + + print("Grid Topology:") + print(" - Fixed 256 buckets") + print(" - CPL-based indexing") + print(" - Explicit MRU peer ordering") + print(" - Better for cross-implementation scenarios\n") + + print("Choice:") + print(" from libp2p.kad_dht import RoutingTable") + print(" rt = RoutingTable(local_id) # Standard k-bucket\n") + + print(" from libp2p.kad_dht import GridRoutingTable") + print(" rt = GridRoutingTable(local_id) # Grid topology\n") + + +def main(): + """Run all examples.""" + print("\n" + "=" * 60) + print("Grid Topology DHT - Comprehensive Example") + print("=" * 60) + + # Run examples + rt, local_id, peer_ids = example_basic_grid_topology() + example_peer_lookup(rt, peer_ids) + example_bucket_distribution(rt) + example_grid_topology_advantages() + example_compare_with_standard_kbucket() + + print("=" * 60) + print("Examples completed!") + print("=" * 60) + + +if __name__ == "__main__": + main() diff --git a/libp2p/kad_dht/__init__.py b/libp2p/kad_dht/__init__.py index 690d37bae..a8f35cf56 100644 --- a/libp2p/kad_dht/__init__.py +++ b/libp2p/kad_dht/__init__.py @@ -14,6 +14,18 @@ from .routing_table import ( RoutingTable, ) +from .grid_routing_table import ( + GridRoutingTable, + GridBucket, + NodeId, +) +from .grid_topology_config import ( + GridTopologyConfig, + get_default_config, + get_testing_config, + get_small_network_config, + get_large_network_config, +) from .utils import ( create_key_from_binary, ) @@ -27,4 +39,12 @@ "PeerRouting", "ValueStore", "create_key_from_binary", + "GridRoutingTable", + "GridBucket", + "NodeId", + "GridTopologyConfig", + "get_default_config", + "get_testing_config", + "get_small_network_config", + "get_large_network_config", ] diff --git a/libp2p/kad_dht/grid_routing_table.py b/libp2p/kad_dht/grid_routing_table.py new file mode 100644 index 000000000..952af7698 --- /dev/null +++ b/libp2p/kad_dht/grid_routing_table.py @@ -0,0 +1,439 @@ +""" +Grid Topology (Kademlia DHT) Routing Table Implementation. + +This implements a 256-bucket binary tree structure based on XOR distance metrics, +matching the cpp-libp2p grid topology implementation. + +Key features: +- 256 fixed buckets (binary tree structure) +- Common Prefix Length (CPL) based bucket indexing +- MRU (Most Recently Used) peer ordering +- Replaceable peer tracking (temporary vs permanent peers) +- Connection status tracking +""" + +from dataclasses import dataclass +import hashlib +import logging +from typing import Any + +import multihash + +from libp2p.peer.id import ID +from libp2p.peer.peerinfo import PeerInfo + +logger = logging.getLogger(__name__) + +GRID_BUCKET_COUNT = 256 +DEFAULT_MAX_BUCKET_SIZE = 20 + + +@dataclass +class BucketPeerInfo: + """Information about a peer in a bucket.""" + + peer_id: ID + peer_info: PeerInfo | None = None + is_replaceable: bool = False + is_connected: bool = False + + +class NodeId: + """DHT Node ID with SHA256-based hashing and XOR distance calculation.""" + + def __init__(self, peer_id: ID): + """Initialize Node ID from a peer ID.""" + self.peer_id: ID | None = peer_id + digest = hashlib.sha256(peer_id.to_bytes()).digest() + mh_bytes = multihash.encode(digest, "sha2-256") + self.data = multihash.decode(mh_bytes).digest + + @classmethod + def from_hash(cls, hash_data: bytes) -> "NodeId": + """Create a NodeId from a pre-computed hash.""" + node_id: NodeId = cls.__new__(cls) # type: ignore[assignment] + node_id.peer_id = None + node_id.data = hash_data + return node_id + + def distance(self, other: "NodeId") -> bytes: + """Calculate XOR distance to another NodeId.""" + distance = bytes(a ^ b for a, b in zip(self.data, other.data)) + return distance + + def common_prefix_len(self, other: "NodeId") -> int: + r""" + Calculate the number of common prefix bits between two node IDs. + + Returns the number of leading bits that are the same. + + Example: + 0x00 and 0xFF have 0 common prefix bits. + 0xFF and 0xFE have 7 common prefix bits. + + """ + distance = self.distance(other) + + for i, byte in enumerate(distance): + if byte != 0: + leading_zeros = 0 + bit = 0x80 + while (byte & bit) == 0: + leading_zeros += 1 + bit >>= 1 + return i * 8 + leading_zeros + + return 256 + + def __eq__(self, other: Any) -> bool: + if isinstance(other, NodeId): + return self.data == other.data + return False + + def __repr__(self) -> str: + return f"NodeId({self.data.hex()[:16]}...)" + + +class GridBucket: + """ + A k-bucket in the grid topology. + + Stores up to k peers, with MRU (Most Recently Used) ordering. + Uses a list to maintain insertion order (LRU at index 0, MRU at end). + """ + + def __init__(self, max_size: int = DEFAULT_MAX_BUCKET_SIZE): + """Initialize a grid bucket.""" + self.max_size = max_size + self.peers: list[BucketPeerInfo] = [] + + def size(self) -> int: + """Get the number of peers in the bucket.""" + return len(self.peers) + + def add_peer( + self, + peer_id: ID, + peer_info: PeerInfo | None = None, + is_replaceable: bool = False, + is_connected: bool = False, + ) -> bool: + """ + Add a peer to the bucket. + + If the peer already exists, update its status and move to end (MRU). + If the bucket is full, return False (caller should handle replacement). + + :param peer_id: ID of the peer to add + :param peer_info: Optional PeerInfo object + :param is_replaceable: True for temporary peers + :param is_connected: True if peer is currently connected + :return: True if peer was added, False if bucket is full + """ + for i, peer_info_obj in enumerate(self.peers): + if peer_info_obj.peer_id == peer_id: + peer_info_obj.is_connected = is_connected + self.peers.append(self.peers.pop(i)) + return True + + if len(self.peers) < self.max_size: + self.peers.append( + BucketPeerInfo( + peer_id=peer_id, + peer_info=peer_info, + is_replaceable=is_replaceable, + is_connected=is_connected, + ) + ) + return True + + return False + + def move_to_front(self, peer_id: ID) -> bool: + """ + Move a peer to the end (most recently used). + + :param peer_id: ID of the peer to move + :return: True if peer was found, False otherwise + """ + for i, peer in enumerate(self.peers): + if peer.peer_id == peer_id: + peer.is_connected = True + self.peers.append(self.peers.pop(i)) + return True + return False + + def remove_replaceable_peer(self) -> ID | None: + """ + Remove a replaceable (temporary) peer from the bucket. + + Searches from end to beginning for the first replaceable unconnected peer. + + :return: ID of removed peer, or None if no replaceable peer found + """ + for i in range(len(self.peers) - 1, -1, -1): + peer = self.peers[i] + if peer.is_replaceable and not peer.is_connected: + removed_id = peer.peer_id + del self.peers[i] + return removed_id + return None + + def remove_peer(self, peer_id: ID) -> bool: + """ + Remove a specific peer from the bucket. + + :param peer_id: ID of the peer to remove + :return: True if peer was removed, False if not found + """ + for i, peer in enumerate(self.peers): + if peer.peer_id == peer_id: + del self.peers[i] + return True + return False + + def contains(self, peer_id: ID) -> bool: + """Check if a peer is in the bucket.""" + return any(peer.peer_id == peer_id for peer in self.peers) + + def get_peer_info(self, peer_id: ID) -> PeerInfo | None: + """Get PeerInfo for a specific peer.""" + for peer in self.peers: + if peer.peer_id == peer_id: + return peer.peer_info + return None + + def peer_ids(self) -> list[ID]: + """Get all peer IDs in the bucket.""" + return [peer.peer_id for peer in self.peers] + + def peer_infos(self) -> list[BucketPeerInfo]: + """Get all BucketPeerInfo objects in the bucket.""" + return list(self.peers) + + def truncate(self, limit: int) -> None: + """Truncate the bucket to a maximum size.""" + while len(self.peers) > limit: + del self.peers[0] + + +class GridRoutingTable: + """ + 256-bucket grid topology routing table for Kademlia DHT. + + Uses a fixed array of 256 buckets indexed by common prefix length (CPL). + Bucket index = 255 - CPL(local_id, peer_id) + """ + + def __init__(self, local_id: ID, max_bucket_size: int = DEFAULT_MAX_BUCKET_SIZE): + """ + Initialize the grid routing table. + + :param local_id: The local peer's ID + :param max_bucket_size: Maximum peers per bucket (default 20) + """ + self.local_id = local_id + self.local_node_id = NodeId(local_id) + self.max_bucket_size = max_bucket_size + + self.buckets: list[GridBucket] = [ + GridBucket(max_bucket_size) for _ in range(GRID_BUCKET_COUNT) + ] + + logger.debug( + f"Initialized grid routing table with {GRID_BUCKET_COUNT} buckets, " + f"max_bucket_size={max_bucket_size}" + ) + + def _get_bucket_index(self, node_id: NodeId) -> int | None: + """ + Calculate the bucket index for a node ID. + + Bucket index = 255 - common_prefix_len(local_id, node_id) + + Returns None if the node ID is the same as local ID. + + :param node_id: The node ID to get bucket for + :return: Bucket index (0-255) or None if node is self + """ + if node_id == self.local_node_id: + return None + + cpl = self.local_node_id.common_prefix_len(node_id) + bucket_index = 255 - cpl + return bucket_index + + def update( + self, + peer_id: ID, + peer_info: PeerInfo | None = None, + is_permanent: bool = True, + is_connected: bool = False, + ) -> bool: + """ + Update or add a peer to the routing table. + + :param peer_id: ID of the peer to add/update + :param peer_info: Optional PeerInfo object + :param is_permanent: True for permanent peers, False for temporary + :param is_connected: True if peer is currently connected + :return: True if added/updated, False if bucket full and no replacement + """ + if peer_id == self.local_id: + return False + + node_id = NodeId(peer_id) + bucket_index = self._get_bucket_index(node_id) + + if bucket_index is None: + return False + + bucket = self.buckets[bucket_index] + + if bucket.add_peer( + peer_id, + peer_info=peer_info, + is_replaceable=not is_permanent, + is_connected=is_connected, + ): + return True + + removed_id = bucket.remove_replaceable_peer() + if removed_id is not None: + bucket.add_peer( + peer_id, + peer_info=peer_info, + is_replaceable=not is_permanent, + is_connected=is_connected, + ) + logger.debug( + f"Replaced peer {removed_id} with {peer_id} in bucket {bucket_index}" + ) + return True + + logger.debug(f"Bucket {bucket_index} full and no replaceable peers") + return False + + def remove(self, peer_id: ID) -> bool: + """ + Remove a peer from the routing table. + + :param peer_id: ID of the peer to remove + :return: True if peer was removed, False if not found + """ + if peer_id == self.local_id: + return False + + node_id = NodeId(peer_id) + bucket_index = self._get_bucket_index(node_id) + + if bucket_index is None: + return False + + return self.buckets[bucket_index].remove_peer(peer_id) + + def get_nearest_peers(self, target_key: bytes, count: int) -> list[ID]: + """ + Find the nearest peers to a target key. + + Implements Kademlia's nearest peer lookup algorithm: + 1. Start with the bucket corresponding to the target key + 2. Expand search to adjacent buckets based on XOR distance + 3. Sort all results by XOR distance + 4. Return top `count` peers + + :param target_key: The target key (bytes) + :param count: Maximum number of peers to return + :return: List of peer IDs, sorted by distance to target key + """ + target_node = NodeId.from_hash(target_key) + + cpl = self.local_node_id.common_prefix_len(target_node) + bucket_index = 255 - cpl + + result_peers: list[tuple[ID, bytes]] = [] # (peer_id, distance) + + def bit_set(distance: bytes, i: int) -> bool: + j = 255 - i + byte_idx = j // 8 + bit_idx = 7 - (j % 8) + return ((distance[byte_idx] >> bit_idx) & 1) != 0 + + target_distance = self.local_node_id.distance(target_node) + + if 0 <= bucket_index < GRID_BUCKET_COUNT: + for peer_info in self.buckets[bucket_index].peer_infos(): + peer_node = NodeId(peer_info.peer_id) + distance = peer_node.distance(target_node) + result_peers.append((peer_info.peer_id, distance)) + + i = bucket_index + while i > 0 and len(result_peers) < count: + i -= 1 + if bit_set(target_distance, i): + for peer_info in self.buckets[i].peer_infos(): + peer_node = NodeId(peer_info.peer_id) + distance = peer_node.distance(target_node) + result_peers.append((peer_info.peer_id, distance)) + + if bucket_index != 0: + for peer_info in self.buckets[0].peer_infos(): + peer_node = NodeId(peer_info.peer_id) + distance = peer_node.distance(target_node) + result_peers.append((peer_info.peer_id, distance)) + + for i in range(1, GRID_BUCKET_COUNT): + if i < bucket_index or (i == bucket_index): + continue + if not bit_set(target_distance, i): + for peer_info in self.buckets[i].peer_infos(): + peer_node = NodeId(peer_info.peer_id) + distance = peer_node.distance(target_node) + result_peers.append((peer_info.peer_id, distance)) + + result_peers.sort(key=lambda x: int.from_bytes(x[1], byteorder="big")) + + return [peer_id for peer_id, _ in result_peers[:count]] + + def get_all_peers(self) -> list[ID]: + """Get all peer IDs in the routing table.""" + peers = [] + for bucket in self.buckets: + peers.extend(bucket.peer_ids()) + return peers + + def contains(self, peer_id: ID) -> bool: + """Check if a peer is in the routing table.""" + if peer_id == self.local_id: + return False + + node_id = NodeId(peer_id) + bucket_index = self._get_bucket_index(node_id) + + if bucket_index is None: + return False + + return self.buckets[bucket_index].contains(peer_id) + + def size(self) -> int: + """Get the total number of peers in the routing table.""" + total = 0 + for bucket in self.buckets: + total += bucket.size() + return total + + def get_bucket(self, index: int) -> GridBucket | None: + """Get a specific bucket by index.""" + if 0 <= index < GRID_BUCKET_COUNT: + return self.buckets[index] + return None + + def get_bucket_stats(self) -> dict[str, Any]: + """Get statistics about bucket distribution.""" + stats = { + "total_peers": self.size(), + "total_buckets": GRID_BUCKET_COUNT, + "non_empty_buckets": sum(1 for b in self.buckets if b.size() > 0), + "bucket_distribution": [b.size() for b in self.buckets], + } + return stats diff --git a/libp2p/kad_dht/grid_topology_config.py b/libp2p/kad_dht/grid_topology_config.py new file mode 100644 index 000000000..b1befce27 --- /dev/null +++ b/libp2p/kad_dht/grid_topology_config.py @@ -0,0 +1,46 @@ +""" +Grid Topology Configuration + +Configuration parameters for the grid topology routing table. +Simplified to include only parameters directly used by GridRoutingTable. +""" + +from dataclasses import dataclass + + +@dataclass +class GridTopologyConfig: + """ + Configuration for grid topology routing table. + + Includes only essential parameters for grid topology peer management. + """ + + # Bucket configuration - directly used by GridRoutingTable + max_bucket_size: int = 20 + """Maximum number of peers per bucket (k-parameter in Kademlia).""" + + def __post_init__(self) -> None: + """Validate configuration after initialization.""" + if self.max_bucket_size < 2: + raise ValueError("max_bucket_size must be at least 2") + + +def get_default_config() -> GridTopologyConfig: + """Get the default grid topology configuration.""" + return GridTopologyConfig(max_bucket_size=20) + + +def get_testing_config() -> GridTopologyConfig: + """Get a configuration suitable for testing with smaller buckets.""" + return GridTopologyConfig(max_bucket_size=5) + + +def get_small_network_config() -> GridTopologyConfig: + """Get a configuration for small networks (local testing).""" + return GridTopologyConfig(max_bucket_size=10) + + +def get_large_network_config() -> GridTopologyConfig: + """Get a configuration for large networks (production).""" + return GridTopologyConfig(max_bucket_size=20) diff --git a/newsfragments/1293.feature.rst b/newsfragments/1293.feature.rst new file mode 100644 index 000000000..cb5cd36f3 --- /dev/null +++ b/newsfragments/1293.feature.rst @@ -0,0 +1 @@ +Added Grid Topology (Kademlia DHT) routing table implementation with 256 fixed buckets and CPL-based indexing. diff --git a/tests/libp2p/kad_dht/test_grid_routing_table.py b/tests/libp2p/kad_dht/test_grid_routing_table.py new file mode 100644 index 000000000..c76e78d59 --- /dev/null +++ b/tests/libp2p/kad_dht/test_grid_routing_table.py @@ -0,0 +1,263 @@ +""" +Tests for the Grid Routing Table implementation. +""" + +import pytest + +from libp2p.kad_dht.grid_routing_table import ( + GRID_BUCKET_COUNT, + GridBucket, + GridRoutingTable, + NodeId, +) +from libp2p.peer.id import ID + + +class TestNodeId: + """Tests for NodeId class.""" + + def test_node_id_creation(self): + """Test creating a NodeId from a peer ID.""" + peer_id = ID.from_base58("QmaCpDMGvV2BGHeYERUEnRQAwe5CcqarqmtA7xNXT92p2") + node_id = NodeId(peer_id) + + assert node_id.data is not None + assert len(node_id.data) == 32 + + def test_node_id_distance(self): + """Test XOR distance calculation between node IDs.""" + peer_id1 = ID.from_base58("QmaCpDMGvV2BGHeYERUEnRQAwe5CcqarqmtA7xNXT92p2") + peer_id2 = ID.from_base58("QmZLaXk3bbiHgVK3zp5A8n2DEuvMZFRv1GAjTrSvZuLnFr") + + node_id1 = NodeId(peer_id1) + node_id2 = NodeId(peer_id2) + + distance = node_id1.distance(node_id2) + assert isinstance(distance, bytes) + assert len(distance) == 32 + + def test_node_id_common_prefix_len(self): + """Test common prefix length calculation.""" + peer_id = ID.from_base58("QmaCpDMGvV2BGHeYERUEnRQAwe5CcqarqmtA7xNXT92p2") + node_id = NodeId(peer_id) + + assert node_id.common_prefix_len(node_id) == 256 + + def test_node_id_equality(self): + """Test NodeId equality.""" + peer_id1 = ID.from_base58("QmaCpDMGvV2BGHeYERUEnRQAwe5CcqarqmtA7xNXT92p2") + node_id1a = NodeId(peer_id1) + node_id1b = NodeId(peer_id1) + + assert node_id1a == node_id1b + + +class TestGridBucket: + """Tests for GridBucket class.""" + + def test_bucket_creation(self): + """Test creating a grid bucket.""" + bucket = GridBucket(max_size=20) + assert bucket.size() == 0 + assert bucket.max_size == 20 + + def test_add_peer_to_empty_bucket(self): + """Test adding a peer to an empty bucket.""" + bucket = GridBucket(max_size=20) + peer_id = ID.from_base58("QmaCpDMGvV2BGHeYERUEnRQAwe5CcqarqmtA7xNXT92p2") + + success = bucket.add_peer(peer_id) + assert success + assert bucket.size() == 1 + assert bucket.contains(peer_id) + + def test_add_duplicate_peer_moves_to_end(self): + """Test adding a duplicate peer moves it to MRU position.""" + bucket = GridBucket(max_size=20) + peer_id1 = ID.from_base58("QmaCpDMGvV2BGHeYERUEnRQAwe5CcqarqmtA7xNXT92p2") + peer_id2 = ID.from_base58("QmZLaXk3bbiHgVK3zp5A8n2DEuvMZFRv1GAjTrSvZuLnFr") + + bucket.add_peer(peer_id1) + bucket.add_peer(peer_id2) + + success = bucket.add_peer(peer_id1) + assert success + assert bucket.size() == 2 + assert list(bucket.peers)[-1].peer_id == peer_id1 + + def test_bucket_full_rejection(self): + """Test that adding to a full bucket returns False.""" + bucket = GridBucket(max_size=2) + peer_id1 = ID.from_base58("QmaCpDMGvV2BGHeYERUEnRQAwe5CcqarqmtA7xNXT92p2") + peer_id2 = ID.from_base58("QmZLaXk3bbiHgVK3zp5A8n2DEuvMZFRv1GAjTrSvZuLnFr") + peer_id3 = ID.from_base58("QmdgB6x6xfBLvV9VwSPj9D7aHCmXhvEVVBn2CUEbNJnTg") + + bucket.add_peer(peer_id1) + bucket.add_peer(peer_id2) + + success = bucket.add_peer(peer_id3) + assert not success + assert bucket.size() == 2 + + def test_remove_replaceable_peer(self): + """Test removing a replaceable (temporary) peer.""" + bucket = GridBucket(max_size=5) + peer_id1 = ID.from_base58("QmaCpDMGvV2BGHeYERUEnRQAwe5CcqarqmtA7xNXT92p2") + peer_id2 = ID.from_base58("QmZLaXk3bbiHgVK3zp5A8n2DEuvMZFRv1GAjTrSvZuLnFr") + + bucket.add_peer(peer_id1, is_replaceable=False) + bucket.add_peer(peer_id2, is_replaceable=True) + + removed = bucket.remove_replaceable_peer() + assert removed == peer_id2 + assert bucket.size() == 1 + assert bucket.contains(peer_id1) + + def test_remove_specific_peer(self): + """Test removing a specific peer.""" + bucket = GridBucket(max_size=20) + peer_id = ID.from_base58("QmaCpDMGvV2BGHeYERUEnRQAwe5CcqarqmtA7xNXT92p2") + + bucket.add_peer(peer_id) + assert bucket.contains(peer_id) + + success = bucket.remove_peer(peer_id) + assert success + assert not bucket.contains(peer_id) + + +class TestGridRoutingTable: + """Tests for GridRoutingTable class.""" + + def test_routing_table_creation(self): + """Test creating a grid routing table.""" + local_id = ID.from_base58("QmaCpDMGvV2BGHeYERUEnRQAwe5CcqarqmtA7xNXT92p2") + rt = GridRoutingTable(local_id) + + assert rt.local_id == local_id + assert len(rt.buckets) == GRID_BUCKET_COUNT + assert rt.size() == 0 + + def test_add_peer_updates_routing_table(self): + """Test adding a peer to the routing table.""" + local_id = ID.from_base58("QmaCpDMGvV2BGHeYERUEnRQAwe5CcqarqmtA7xNXT92p2") + peer_id = ID.from_base58("QmZLaXk3bbiHgVK3zp5A8n2DEuvMZFRv1GAjTrSvZuLnFr") + + rt = GridRoutingTable(local_id) + success = rt.update(peer_id) + + assert success + assert rt.size() == 1 + assert rt.contains(peer_id) + + def test_cannot_add_self(self): + """Test that local peer cannot be added to routing table.""" + local_id = ID.from_base58("QmaCpDMGvV2BGHeYERUEnRQAwe5CcqarqmtA7xNXT92p2") + rt = GridRoutingTable(local_id) + + success = rt.update(local_id) + assert not success + assert rt.size() == 0 + + def test_remove_peer_from_routing_table(self): + """Test removing a peer from the routing table.""" + local_id = ID.from_base58("QmaCpDMGvV2BGHeYERUEnRQAwe5CcqarqmtA7xNXT92p2") + peer_id = ID.from_base58("QmZLaXk3bbiHgVK3zp5A8n2DEuvMZFRv1GAjTrSvZuLnFr") + + rt = GridRoutingTable(local_id) + rt.update(peer_id) + assert rt.contains(peer_id) + + success = rt.remove(peer_id) + assert success + assert not rt.contains(peer_id) + + def test_get_all_peers(self): + """Test getting all peers from routing table.""" + local_id = ID.from_base58("QmaCpDMGvV2BGHeYERUEnRQAwe5CcqarqmtA7xNXT92p2") + peer_id1 = ID.from_base58("QmZLaXk3bbiHgVK3zp5A8n2DEuvMZFRv1GAjTrSvZuLnFr") + peer_id2 = ID.from_base58("QmdgB6x6xfBLvV9VwSPj9D7aHCmXhvEVVBn2CUEbNJnTg") + + rt = GridRoutingTable(local_id) + rt.update(peer_id1) + rt.update(peer_id2) + + all_peers = rt.get_all_peers() + assert len(all_peers) == 2 + assert peer_id1 in all_peers + assert peer_id2 in all_peers + + def test_get_bucket_index(self): + """Test getting bucket index for a peer.""" + local_id = ID.from_base58("QmaCpDMGvV2BGHeYERUEnRQAwe5CcqarqmtA7xNXT92p2") + peer_id = ID.from_base58("QmZLaXk3bbiHgVK3zp5A8n2DEuvMZFRv1GAjTrSvZuLnFr") + + rt = GridRoutingTable(local_id) + peer_node = NodeId(peer_id) + + bucket_index = rt._get_bucket_index(peer_node) + assert bucket_index is not None + assert 0 <= bucket_index < GRID_BUCKET_COUNT + + def test_bucket_index_none_for_self(self): + """Test that bucket index is None for local ID.""" + local_id = ID.from_base58("QmaCpDMGvV2BGHeYERUEnRQAwe5CcqarqmtA7xNXT92p2") + rt = GridRoutingTable(local_id) + + bucket_index = rt._get_bucket_index(rt.local_node_id) + assert bucket_index is None + + def test_multiple_peers_in_same_bucket(self): + """Test adding multiple peers to the same bucket.""" + local_id = ID.from_base58("QmaCpDMGvV2BGHeYERUEnRQAwe5CcqarqmtA7xNXT92p2") + + rt = GridRoutingTable(local_id) + + peers = [ + ID.from_base58("QmZLaXk3bbiHgVK3zp5A8n2DEuvMZFRv1GAjTrSvZuLnFr"), + ID.from_base58("QmdgB6x6xfBLvV9VwSPj9D7aHCmXhvEVVBn2CUEbNJnTg"), + ID.from_base58("QmYPp4CUQpRcpHBqWgmQ4TjqW9AHJ3qLhQV5d7s3d8hMJx"), + ] + + for peer_id in peers: + rt.update(peer_id) + + assert rt.size() == len(peers) + for peer_id in peers: + assert rt.contains(peer_id) + + def test_get_bucket_stats(self): + """Test getting bucket statistics.""" + local_id = ID.from_base58("QmaCpDMGvV2BGHeYERUEnRQAwe5CcqarqmtA7xNXT92p2") + peer_id = ID.from_base58("QmZLaXk3bbiHgVK3zp5A8n2DEuvMZFRv1GAjTrSvZuLnFr") + + rt = GridRoutingTable(local_id) + rt.update(peer_id) + + stats = rt.get_bucket_stats() + assert stats["total_peers"] == 1 + assert stats["total_buckets"] == GRID_BUCKET_COUNT + assert stats["non_empty_buckets"] == 1 + + def test_permanent_vs_replaceable_peers(self): + """Test that replaceable peers can be replaced.""" + local_id = ID.from_base58("QmaCpDMGvV2BGHeYERUEnRQAwe5CcqarqmtA7xNXT92p2") + + rt = GridRoutingTable(local_id, max_bucket_size=2) + + peers = [ + ID.from_base58("QmZLaXk3bbiHgVK3zp5A8n2DEuvMZFRv1GAjTrSvZuLnFr"), + ID.from_base58("QmdgB6x6xfBLvV9VwSPj9D7aHCmXhvEVVBn2CUEbNJnTg"), + ID.from_base58("QmYPp4CUQpRcpHBqWgmQ4TjqW9AHJ3qLhQV5d7s3d8hMJx"), + ] + + rt.update(peers[0], is_permanent=False) + rt.update(peers[1], is_permanent=True) + + rt.update(peers[2], is_permanent=False) + assert rt.size() >= 1 + assert rt.size() <= 3 + + +if __name__ == "__main__": + pytest.main([__file__, "-v"])