|
17 | 17 | use std::collections::HashMap;
|
18 | 18 |
|
19 | 19 | use stacks_common::deps_common::bitcoin::network::serialize::BitcoinHash;
|
| 20 | +use stacks_common::types::net::PeerAddress; |
| 21 | +use stacks_common::util::hash::Hash160; |
| 22 | +use stacks_common::util::secp256k1::{Secp256k1PrivateKey, Secp256k1PublicKey}; |
20 | 23 |
|
21 | 24 | use crate::burnchains::bitcoin::indexer::BitcoinIndexer;
|
22 | 25 | use crate::burnchains::db::BurnchainHeaderReader;
|
@@ -2014,3 +2017,231 @@ fn test_sync_inv_2_peers_different_pox_vectors() {
|
2014 | 2017 | assert!(!peer_2_inv.has_ith_microblock_stream(num_blocks - 2 * reward_cycle_length));
|
2015 | 2018 | })
|
2016 | 2019 | }
|
| 2020 | + |
| 2021 | +/// Helper function to create a test neighbor without binding to a port |
| 2022 | +fn create_test_neighbor(port: u16) -> Neighbor { |
| 2023 | + let private_key = Secp256k1PrivateKey::random(); |
| 2024 | + let public_key = Secp256k1PublicKey::from_private(&private_key); |
| 2025 | + let public_key_hash = Hash160::from_node_public_key(&public_key); |
| 2026 | + |
| 2027 | + let neighbor_address = NeighborAddress { |
| 2028 | + addrbytes: PeerAddress::from_ipv4(127, 0, 0, 1), |
| 2029 | + port, |
| 2030 | + public_key_hash, |
| 2031 | + }; |
| 2032 | + |
| 2033 | + let neighbor_key = NeighborKey::from_neighbor_address( |
| 2034 | + 0x18000000, // peer_version |
| 2035 | + 0x80000000, // network_id |
| 2036 | + &neighbor_address, |
| 2037 | + ); |
| 2038 | + |
| 2039 | + Neighbor::empty(&neighbor_key, &public_key, 9999999) |
| 2040 | +} |
| 2041 | + |
| 2042 | +/// Specification for a test peer |
| 2043 | +struct TestPeerSpec { |
| 2044 | + port: u16, |
| 2045 | + status: NodeStatus, |
| 2046 | + num_sortitions: u64, |
| 2047 | +} |
| 2048 | + |
| 2049 | +impl TestPeerSpec { |
| 2050 | + fn new(port: u16, status: NodeStatus, num_sortitions: u64) -> Self { |
| 2051 | + Self { |
| 2052 | + port, |
| 2053 | + status, |
| 2054 | + num_sortitions, |
| 2055 | + } |
| 2056 | + } |
| 2057 | +} |
| 2058 | + |
| 2059 | +/// Helper function to setup multiple peers at once |
| 2060 | +fn setup_test_inv_state(specs: &[TestPeerSpec]) -> (InvState, Vec<Neighbor>) { |
| 2061 | + // first_block_height=100, so heights = 100 + num_sortitions |
| 2062 | + let mut inv_state = InvState::new(100, 60, 3); |
| 2063 | + let mut neighbors = Vec::new(); |
| 2064 | + |
| 2065 | + for spec in specs { |
| 2066 | + let neighbor = create_test_neighbor(spec.port); |
| 2067 | + let neighbor_key = neighbor.addr.clone(); |
| 2068 | + |
| 2069 | + inv_state.add_peer(neighbor_key.clone(), false); |
| 2070 | + if let Some(stats) = inv_state.get_stats_mut(&neighbor_key) { |
| 2071 | + stats.status = spec.status; |
| 2072 | + stats.inv.num_sortitions = spec.num_sortitions; |
| 2073 | + } |
| 2074 | + |
| 2075 | + neighbors.push(neighbor); |
| 2076 | + } |
| 2077 | + |
| 2078 | + (inv_state, neighbors) |
| 2079 | +} |
| 2080 | + |
| 2081 | +/// Helper function to create neighbors without adding them to inv_state |
| 2082 | +fn create_test_neighbors(ports: &[u16]) -> Vec<Neighbor> { |
| 2083 | + ports |
| 2084 | + .iter() |
| 2085 | + .map(|&port| create_test_neighbor(port)) |
| 2086 | + .collect() |
| 2087 | +} |
| 2088 | + |
| 2089 | +/// Helper function to assert the expected heights for both IBD and non-IBD modes |
| 2090 | +fn assert_max_heights( |
| 2091 | + inv_state: &InvState, |
| 2092 | + neighbors: &[Neighbor], |
| 2093 | + expected_non_ibd: Option<u64>, |
| 2094 | + expected_ibd: Option<u64>, |
| 2095 | +) { |
| 2096 | + assert_eq!( |
| 2097 | + inv_state.get_max_stacks_height_of_neighbors(neighbors, false), |
| 2098 | + expected_non_ibd, |
| 2099 | + "Non-IBD mode assertion failed" |
| 2100 | + ); |
| 2101 | + assert_eq!( |
| 2102 | + inv_state.get_max_stacks_height_of_neighbors(neighbors, true), |
| 2103 | + expected_ibd, |
| 2104 | + "IBD mode assertion failed" |
| 2105 | + ); |
| 2106 | +} |
| 2107 | + |
| 2108 | +#[test] |
| 2109 | +fn test_get_max_stacks_height_of_neighbors() { |
| 2110 | + // Test empty neighbors list |
| 2111 | + let (inv_state, neighbors) = setup_test_inv_state(&[]); |
| 2112 | + assert_max_heights(&inv_state, &neighbors, None, None); |
| 2113 | +} |
| 2114 | + |
| 2115 | +#[test] |
| 2116 | +fn test_get_max_stacks_height_of_neighbors_no_stats() { |
| 2117 | + // Test neighbors without any stats in the inv_state |
| 2118 | + let inv_state = InvState::new(100, 60, 3); |
| 2119 | + let neighbors = create_test_neighbors(&[8080]); |
| 2120 | + |
| 2121 | + // Should return None since no stats exist for the neighbor |
| 2122 | + assert_max_heights(&inv_state, &neighbors, None, None); |
| 2123 | +} |
| 2124 | + |
| 2125 | +#[test] |
| 2126 | +fn test_get_max_stacks_height_of_neighbors_single_online_peer() { |
| 2127 | + let (inv_state, neighbors) = |
| 2128 | + setup_test_inv_state(&[TestPeerSpec::new(8080, NodeStatus::Online, 150)]); |
| 2129 | + |
| 2130 | + // Online peers should be accepted in both modes (height = 100 + 150 = 250) |
| 2131 | + assert_max_heights(&inv_state, &neighbors, Some(250), Some(250)); |
| 2132 | +} |
| 2133 | + |
| 2134 | +#[test] |
| 2135 | +fn test_get_max_stacks_height_of_neighbors_single_diverged_peer() { |
| 2136 | + let (inv_state, neighbors) = |
| 2137 | + setup_test_inv_state(&[TestPeerSpec::new(8080, NodeStatus::Diverged, 200)]); |
| 2138 | + |
| 2139 | + // Diverged peers accepted only in IBD mode (height = 100 + 200 = 300) |
| 2140 | + assert_max_heights(&inv_state, &neighbors, None, Some(300)); |
| 2141 | +} |
| 2142 | + |
| 2143 | +#[test] |
| 2144 | +fn test_get_max_stacks_height_of_neighbors_single_broken_peer() { |
| 2145 | + let (inv_state, neighbors) = |
| 2146 | + setup_test_inv_state(&[TestPeerSpec::new(8080, NodeStatus::Broken, 180)]); |
| 2147 | + |
| 2148 | + // Broken peers never accepted |
| 2149 | + assert_max_heights(&inv_state, &neighbors, None, None); |
| 2150 | +} |
| 2151 | + |
| 2152 | +#[test] |
| 2153 | +fn test_get_max_stacks_height_of_neighbors_single_stale_peer() { |
| 2154 | + let (inv_state, neighbors) = |
| 2155 | + setup_test_inv_state(&[TestPeerSpec::new(8080, NodeStatus::Stale, 120)]); |
| 2156 | + |
| 2157 | + // Stale peers never accepted |
| 2158 | + assert_max_heights(&inv_state, &neighbors, None, None); |
| 2159 | +} |
| 2160 | + |
| 2161 | +#[test] |
| 2162 | +fn test_get_max_stacks_height_of_neighbors_single_dead_peer() { |
| 2163 | + let (inv_state, neighbors) = |
| 2164 | + setup_test_inv_state(&[TestPeerSpec::new(8080, NodeStatus::Dead, 190)]); |
| 2165 | + |
| 2166 | + // Dead peers never accepted |
| 2167 | + assert_max_heights(&inv_state, &neighbors, None, None); |
| 2168 | +} |
| 2169 | + |
| 2170 | +#[test] |
| 2171 | +fn test_get_max_stacks_height_of_neighbors_multiple_online_peers() { |
| 2172 | + let specs = [ |
| 2173 | + TestPeerSpec::new(8080, NodeStatus::Online, 100), // height = 200 |
| 2174 | + TestPeerSpec::new(8081, NodeStatus::Online, 250), // height = 350 (max) |
| 2175 | + TestPeerSpec::new(8082, NodeStatus::Online, 180), // height = 280 |
| 2176 | + ]; |
| 2177 | + |
| 2178 | + let (inv_state, neighbors) = setup_test_inv_state(&specs); |
| 2179 | + |
| 2180 | + // Should return max height among all online peers (350) |
| 2181 | + assert_max_heights(&inv_state, &neighbors, Some(350), Some(350)); |
| 2182 | +} |
| 2183 | + |
| 2184 | +#[test] |
| 2185 | +fn test_get_max_stacks_height_of_neighbors_mixed_statuses_non_ibd() { |
| 2186 | + let specs = [ |
| 2187 | + TestPeerSpec::new(8080, NodeStatus::Online, 150), // height = 250 |
| 2188 | + TestPeerSpec::new(8081, NodeStatus::Diverged, 300), // height = 400 (ignored in non-IBD) |
| 2189 | + TestPeerSpec::new(8082, NodeStatus::Online, 200), // height = 300 (max among Online) |
| 2190 | + TestPeerSpec::new(8083, NodeStatus::Broken, 350), // ignored |
| 2191 | + ]; |
| 2192 | + |
| 2193 | + let (inv_state, neighbors) = setup_test_inv_state(&specs); |
| 2194 | + |
| 2195 | + // Non-IBD: only Online peers considered, max = 300 |
| 2196 | + assert_max_heights(&inv_state, &neighbors, Some(300), Some(400)); |
| 2197 | +} |
| 2198 | + |
| 2199 | +#[test] |
| 2200 | +fn test_get_max_stacks_height_of_neighbors_mixed_statuses_ibd() { |
| 2201 | + let specs = [ |
| 2202 | + TestPeerSpec::new(8080, NodeStatus::Online, 150), // height = 250 |
| 2203 | + TestPeerSpec::new(8081, NodeStatus::Diverged, 300), // height = 400 (max in IBD) |
| 2204 | + TestPeerSpec::new(8082, NodeStatus::Online, 200), // height = 300 |
| 2205 | + TestPeerSpec::new(8083, NodeStatus::Broken, 350), // ignored even in IBD |
| 2206 | + ]; |
| 2207 | + |
| 2208 | + let (inv_state, neighbors) = setup_test_inv_state(&specs); |
| 2209 | + |
| 2210 | + // IBD: both Online and Diverged considered, max = 400 |
| 2211 | + assert_max_heights(&inv_state, &neighbors, Some(300), Some(400)); |
| 2212 | +} |
| 2213 | + |
| 2214 | +#[test] |
| 2215 | +fn test_get_max_stacks_height_of_neighbors_minimum_height() { |
| 2216 | + let (inv_state, neighbors) = |
| 2217 | + setup_test_inv_state(&[TestPeerSpec::new(8080, NodeStatus::Online, 0)]); |
| 2218 | + |
| 2219 | + // Should handle minimum height correctly (height = 100 + 0 = 100) |
| 2220 | + assert_max_heights(&inv_state, &neighbors, Some(100), Some(100)); |
| 2221 | +} |
| 2222 | + |
| 2223 | +#[test] |
| 2224 | +fn test_get_max_stacks_height_of_neighbors_all_invalid_statuses() { |
| 2225 | + let specs = [ |
| 2226 | + TestPeerSpec::new(8080, NodeStatus::Broken, 150), |
| 2227 | + TestPeerSpec::new(8081, NodeStatus::Dead, 200), |
| 2228 | + ]; |
| 2229 | + |
| 2230 | + let (inv_state, neighbors) = setup_test_inv_state(&specs); |
| 2231 | + |
| 2232 | + // All invalid statuses should return None |
| 2233 | + assert_max_heights(&inv_state, &neighbors, None, None); |
| 2234 | +} |
| 2235 | + |
| 2236 | +#[test] |
| 2237 | +fn test_get_max_stacks_height_of_neighbors_diverged_only_non_ibd() { |
| 2238 | + let specs = [ |
| 2239 | + TestPeerSpec::new(8080, NodeStatus::Diverged, 150), // height = 250 |
| 2240 | + TestPeerSpec::new(8081, NodeStatus::Diverged, 200), // height = 300 (max) |
| 2241 | + ]; |
| 2242 | + |
| 2243 | + let (inv_state, neighbors) = setup_test_inv_state(&specs); |
| 2244 | + |
| 2245 | + // Non-IBD: Diverged ignored, IBD: Diverged accepted |
| 2246 | + assert_max_heights(&inv_state, &neighbors, None, Some(300)); |
| 2247 | +} |
0 commit comments