A Python port of the Invisible Internet Project (I2P)
at main 231 lines 8.3 kB view raw
1"""Tests for PeerSelector (tunnel hop selection) and TunnelHopConfig. 2 3TDD: tests written before implementation. 4""" 5 6import os 7import struct 8 9import pytest 10 11from i2p_peer.hop_config import TunnelHopConfig 12from i2p_peer.organizer import PeerTier, ProfileOrganizer 13from i2p_peer.profile import PeerProfile 14from i2p_peer.selector import PeerSelector 15 16 17# --- Helpers --- 18 19def _make_hash(i: int) -> bytes: 20 """Create a unique 32-byte peer hash from an integer.""" 21 return i.to_bytes(4, "big") + b"\x00" * 28 22 23 24def _make_profile(i: int, capacity: float = 0.5, speed: float = 0.5, 25 tunnel_builds: int = 20) -> PeerProfile: 26 """Create a PeerProfile with given scores.""" 27 p = PeerProfile(_make_hash(i)) 28 p.capacity = capacity 29 p.speed = speed 30 p.tunnel_builds_succeeded = tunnel_builds 31 p.tunnel_builds_failed = 0 32 p.record_latency(200.0) # give it some latency data 33 p.heard_from() 34 return p 35 36 37def _organizer_with_profiles(profiles: list[PeerProfile]) -> ProfileOrganizer: 38 org = ProfileOrganizer() 39 for p in profiles: 40 org.add_profile(p) 41 return org 42 43 44# --- TunnelHopConfig tests --- 45 46class TestTunnelHopConfig: 47 def test_fields_present(self): 48 cfg = TunnelHopConfig( 49 peer_hash=os.urandom(32), 50 receive_tunnel_id=12345, 51 layer_key=os.urandom(32), 52 iv_key=os.urandom(32), 53 reply_key=os.urandom(32), 54 reply_iv=os.urandom(16), 55 is_gateway=True, 56 is_endpoint=False, 57 ) 58 assert len(cfg.peer_hash) == 32 59 assert len(cfg.layer_key) == 32 60 assert len(cfg.iv_key) == 32 61 assert len(cfg.reply_key) == 32 62 assert len(cfg.reply_iv) == 16 63 assert cfg.receive_tunnel_id == 12345 64 assert cfg.is_gateway is True 65 assert cfg.is_endpoint is False 66 67 def test_defaults(self): 68 cfg = TunnelHopConfig( 69 peer_hash=os.urandom(32), 70 receive_tunnel_id=1, 71 layer_key=os.urandom(32), 72 iv_key=os.urandom(32), 73 reply_key=os.urandom(32), 74 reply_iv=os.urandom(16), 75 ) 76 assert cfg.is_gateway is False 77 assert cfg.is_endpoint is False 78 79 80# --- PeerSelector.select_peers tests --- 81 82class TestSelectPeers: 83 def test_basic_selection(self): 84 """Select N peers from profiles.""" 85 profiles = [_make_profile(i, capacity=0.5) for i in range(10)] 86 org = _organizer_with_profiles(profiles) 87 sel = PeerSelector(org) 88 result = sel.select_peers(3) 89 assert len(result) == 3 90 # All returned items should be 32-byte hashes 91 for h in result: 92 assert isinstance(h, bytes) 93 assert len(h) == 32 94 95 def test_exclude_list_respected(self): 96 """Excluded peers must not appear in results.""" 97 profiles = [_make_profile(i, capacity=0.5) for i in range(5)] 98 org = _organizer_with_profiles(profiles) 99 sel = PeerSelector(org) 100 excluded = {_make_hash(0), _make_hash(1), _make_hash(2)} 101 result = sel.select_peers(5, exclude=excluded) 102 for h in result: 103 assert h not in excluded 104 # Only 2 non-excluded peers exist 105 assert len(result) == 2 106 107 def test_higher_scoring_peers_preferred(self): 108 """Peers with higher capacity/speed should be selected first.""" 109 # Create peers with distinct tiers 110 fast_peer = _make_profile(1, capacity=0.9, speed=0.8, tunnel_builds=20) 111 high_cap_peer = _make_profile(2, capacity=0.7, speed=0.3, tunnel_builds=5) 112 standard_peer = _make_profile(3, capacity=0.4, speed=0.2, tunnel_builds=5) 113 failing_peer = _make_profile(4, capacity=0.1, speed=0.1, tunnel_builds=5) 114 115 org = _organizer_with_profiles([fast_peer, high_cap_peer, standard_peer, failing_peer]) 116 sel = PeerSelector(org) 117 118 # Request 2 peers — should get the best tiers 119 result = sel.select_peers(2) 120 assert len(result) == 2 121 # The fast and high-cap peers should be selected 122 assert _make_hash(1) in result 123 assert _make_hash(2) in result 124 125 def test_empty_peer_set(self): 126 """Empty organizer returns empty list.""" 127 org = ProfileOrganizer() 128 sel = PeerSelector(org) 129 result = sel.select_peers(3) 130 assert result == [] 131 132 def test_fewer_peers_than_requested(self): 133 """If fewer peers exist than requested, return all available.""" 134 profiles = [_make_profile(i, capacity=0.5) for i in range(2)] 135 org = _organizer_with_profiles(profiles) 136 sel = PeerSelector(org) 137 result = sel.select_peers(10) 138 assert len(result) == 2 139 140 141# --- PeerSelector.select_hops tests --- 142 143class TestSelectHops: 144 def test_returns_correct_number(self): 145 """select_hops returns the requested number of TunnelHopConfig objects.""" 146 profiles = [_make_profile(i, capacity=0.5) for i in range(5)] 147 org = _organizer_with_profiles(profiles) 148 sel = PeerSelector(org) 149 hops = sel.select_hops(3) 150 assert len(hops) == 3 151 for hop in hops: 152 assert isinstance(hop, TunnelHopConfig) 153 154 def test_hop_key_sizes(self): 155 """Each hop has correct key sizes.""" 156 profiles = [_make_profile(i, capacity=0.5) for i in range(5)] 157 org = _organizer_with_profiles(profiles) 158 sel = PeerSelector(org) 159 hops = sel.select_hops(3) 160 for hop in hops: 161 assert len(hop.peer_hash) == 32 162 assert len(hop.layer_key) == 32 163 assert len(hop.iv_key) == 32 164 assert len(hop.reply_key) == 32 165 assert len(hop.reply_iv) == 16 166 assert isinstance(hop.receive_tunnel_id, int) 167 assert hop.receive_tunnel_id > 0 168 169 def test_gateway_and_endpoint_flags(self): 170 """First hop is gateway, last hop is endpoint.""" 171 profiles = [_make_profile(i, capacity=0.5) for i in range(5)] 172 org = _organizer_with_profiles(profiles) 173 sel = PeerSelector(org) 174 hops = sel.select_hops(3) 175 176 assert hops[0].is_gateway is True 177 assert hops[0].is_endpoint is False 178 179 assert hops[1].is_gateway is False 180 assert hops[1].is_endpoint is False 181 182 assert hops[2].is_gateway is False 183 assert hops[2].is_endpoint is True 184 185 def test_single_hop_is_both_gateway_and_endpoint(self): 186 """A single-hop tunnel is both gateway and endpoint.""" 187 profiles = [_make_profile(0, capacity=0.5)] 188 org = _organizer_with_profiles(profiles) 189 sel = PeerSelector(org) 190 hops = sel.select_hops(1) 191 assert len(hops) == 1 192 assert hops[0].is_gateway is True 193 assert hops[0].is_endpoint is True 194 195 def test_keys_are_unique_per_hop(self): 196 """Each hop should have distinct random keys.""" 197 profiles = [_make_profile(i, capacity=0.5) for i in range(5)] 198 org = _organizer_with_profiles(profiles) 199 sel = PeerSelector(org) 200 hops = sel.select_hops(3) 201 layer_keys = [hop.layer_key for hop in hops] 202 assert len(set(layer_keys)) == 3 # all unique 203 204 def test_exclude_in_select_hops(self): 205 """Exclude set is passed through to peer selection.""" 206 profiles = [_make_profile(i, capacity=0.5) for i in range(5)] 207 org = _organizer_with_profiles(profiles) 208 sel = PeerSelector(org) 209 excluded = {_make_hash(0), _make_hash(1)} 210 hops = sel.select_hops(3, exclude=excluded) 211 hop_hashes = {hop.peer_hash for hop in hops} 212 for h in excluded: 213 assert h not in hop_hashes 214 215 def test_empty_peers_returns_empty_hops(self): 216 """Empty organizer returns empty hop list.""" 217 org = ProfileOrganizer() 218 sel = PeerSelector(org) 219 hops = sel.select_hops(3) 220 assert hops == [] 221 222 def test_fewer_peers_than_hops(self): 223 """If fewer peers than requested hops, return what's available.""" 224 profiles = [_make_profile(i, capacity=0.5) for i in range(2)] 225 org = _organizer_with_profiles(profiles) 226 sel = PeerSelector(org) 227 hops = sel.select_hops(5) 228 assert len(hops) == 2 229 # First is gateway, last is endpoint 230 assert hops[0].is_gateway is True 231 assert hops[-1].is_endpoint is True