f expand test

Author: TheBlueMatt

lightning/src/ln/peer_handler.rs

@@ -3728,6 +3728,7 @@ mod tests {
 		chan_handler: test_utils::TestChannelMessageHandler,
 		routing_handler: test_utils::TestRoutingMessageHandler,
 		custom_handler: TestCustomMessageHandler,
+		send_only_handler: TestBaseMsgHandler,
 		logger: test_utils::TestLogger,
 		node_signer: test_utils::TestNodeSigner,
 	}
@@ -3780,6 +3781,32 @@ mod tests {
 		}
 	}
 
+	struct TestBaseMsgHandler(test_utils::ConnectionTracker);
+
+	impl BaseMessageHandler for TestBaseMsgHandler {
+		fn get_and_clear_pending_msg_events(&self) -> Vec<MessageSendEvent> {
+			Vec::new()
+		}
+
+		fn peer_disconnected(&self, their_node_id: PublicKey) {
+			self.0.peer_disconnected(their_node_id);
+		}
+
+		fn peer_connected(
+			&self, their_node_id: PublicKey, _msg: &Init, _inbound: bool,
+		) -> Result<(), ()> {
+			self.0.peer_connected(their_node_id)
+		}
+
+		fn provided_node_features(&self) -> NodeFeatures {
+			NodeFeatures::empty()
+		}
+
+		fn provided_init_features(&self, _: PublicKey) -> InitFeatures {
+			InitFeatures::empty()
+		}
+	}
+
 	fn create_peermgr_cfgs(peer_count: usize) -> Vec<PeerManagerCfg> {
 		let mut cfgs = Vec::new();
 		for i in 0..peer_count {
@@ -3796,6 +3823,7 @@ mod tests {
 				logger: test_utils::TestLogger::with_id(i.to_string()),
 				routing_handler: test_utils::TestRoutingMessageHandler::new(),
 				custom_handler: TestCustomMessageHandler::new(features),
+				send_only_handler: TestBaseMsgHandler(test_utils::ConnectionTracker::new()),
 				node_signer: test_utils::TestNodeSigner::new(node_secret),
 			});
 		}
@@ -3819,6 +3847,7 @@ mod tests {
 				logger: test_utils::TestLogger::new(),
 				routing_handler: test_utils::TestRoutingMessageHandler::new(),
 				custom_handler: TestCustomMessageHandler::new(features),
+				send_only_handler: TestBaseMsgHandler(test_utils::ConnectionTracker::new()),
 				node_signer: test_utils::TestNodeSigner::new(node_secret),
 			});
 		}
@@ -3837,6 +3866,7 @@ mod tests {
 				logger: test_utils::TestLogger::new(),
 				routing_handler: test_utils::TestRoutingMessageHandler::new(),
 				custom_handler: TestCustomMessageHandler::new(features),
+				send_only_handler: TestBaseMsgHandler(test_utils::ConnectionTracker::new()),
 				node_signer: test_utils::TestNodeSigner::new(node_secret),
 			});
 		}
@@ -3853,7 +3883,7 @@ mod tests {
 				route_handler: &cfgs[i].routing_handler,
 				onion_message_handler: IgnoringMessageHandler {},
 				custom_message_handler: &cfgs[i].custom_handler,
-				send_only_message_handler: IgnoringMessageHandler {},
+				send_only_message_handler: &cfgs[i].send_only_handler,
 			};
 			let peer = PeerManager::new(
 				msg_handler,
@@ -3876,7 +3906,7 @@ mod tests {
 		&'a test_utils::TestLogger,
 		&'a TestCustomMessageHandler,
 		&'a test_utils::TestNodeSigner,
-		IgnoringMessageHandler,
+		&'a TestBaseMsgHandler,
 	>;
 
 	fn try_establish_connection<'a>(
@@ -4252,6 +4282,7 @@ mod tests {
 		let chan_handler = peers[handler & 1].message_handler.chan_handler;
 		let route_handler = peers[handler & 1].message_handler.route_handler;
 		let custom_message_handler = peers[handler & 1].message_handler.custom_message_handler;
+		let send_only_msg_handler = peers[handler & 1].message_handler.send_only_message_handler;
 
 		match handler & !1 {
 			0 => {
@@ -4263,6 +4294,9 @@ mod tests {
 			4 => {
 				custom_message_handler.conn_tracker.fail_connections.store(true, Ordering::Release);
 			},
+			6 => {
+				send_only_msg_handler.0.fail_connections.store(true, Ordering::Release);
+			},
 			_ => panic!(),
 		}
 		let (_sd1, _sd2, a_refused, b_refused) = try_establish_connection(&peers[0], &peers[1]);
@@ -4278,16 +4312,18 @@ mod tests {
 			chan_handler.conn_tracker.had_peers.load(Ordering::Acquire)
 				|| route_handler.conn_tracker.had_peers.load(Ordering::Acquire)
 				|| custom_message_handler.conn_tracker.had_peers.load(Ordering::Acquire)
+				|| send_only_msg_handler.0.had_peers.load(Ordering::Acquire)
 		);
 		// And both message handlers doing tracking should see the disconnection
 		assert!(chan_handler.conn_tracker.connected_peers.lock().unwrap().is_empty());
 		assert!(route_handler.conn_tracker.connected_peers.lock().unwrap().is_empty());
 		assert!(custom_message_handler.conn_tracker.connected_peers.lock().unwrap().is_empty());
+		assert!(send_only_msg_handler.0.connected_peers.lock().unwrap().is_empty());
 	}
 
 	#[test]
 	fn test_peer_connected_error_disconnects() {
-		for i in 0..6 {
+		for i in 0..8 {
 			do_test_peer_connected_error_disconnects(i);
 		}
 	}