A Python port of the Invisible Internet Project (I2P)
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