Refactor: Update TestMessageRouter to enum

To allow choosing different message router types for testing nodes,
convert `TestMessageRouter` to an enum with variants `DefaultMessageRouter`
and `NodeIdMessageRouter`. This provides better flexibility when testing
various scenarios.

Author: shaavan

lightning/src/ln/channelmanager.rs

@@ -17230,7 +17230,7 @@ pub mod bench {
 		let scorer = RwLock::new(test_utils::TestScorer::new());
 		let entropy = test_utils::TestKeysInterface::new(&[0u8; 32], network);
 		let router = test_utils::TestRouter::new(Arc::new(NetworkGraph::new(network, &logger_a)), &logger_a, &scorer);
-		let message_router = test_utils::TestMessageRouter::new(Arc::new(NetworkGraph::new(network, &logger_a)), &entropy);
+		let message_router = test_utils::TestMessageRouter::new_default(Arc::new(NetworkGraph::new(network, &logger_a)), &entropy);
 
 		let mut config: UserConfig = Default::default();
 		config.channel_config.max_dust_htlc_exposure = MaxDustHTLCExposure::FeeRateMultiplier(5_000_000 / 253);

lightning/src/ln/functional_test_utils.rs

@@ -47,7 +47,7 @@ use crate::util::scid_utils;
 use crate::util::ser::{ReadableArgs, Writeable};
 use crate::util::test_channel_signer::SignerOp;
 use crate::util::test_channel_signer::TestChannelSigner;
-use crate::util::test_utils;
+use crate::util::test_utils::{self, TestLogger};
 use crate::util::test_utils::{TestChainMonitor, TestKeysInterface, TestScorer};
 
 use bitcoin::amount::Amount;
@@ -855,7 +855,7 @@ impl<'a, 'b, 'c> Drop for Node<'a, 'b, 'c> {
 							&self.logger,
 							&scorer,
 						),
-						message_router: &test_utils::TestMessageRouter::new(
+						message_router: &test_utils::TestMessageRouter::new_default(
 							network_graph,
 							self.keys_manager,
 						),
@@ -4204,49 +4204,47 @@ pub fn create_chanmon_cfgs_with_keys(
 	chan_mon_cfgs
 }
 
-pub fn create_node_cfgs<'a>(
-	node_count: usize, chanmon_cfgs: &'a Vec<TestChanMonCfg>,
-) -> Vec<NodeCfg<'a>> {
-	create_node_cfgs_with_persisters(
-		node_count,
-		chanmon_cfgs,
-		chanmon_cfgs.iter().map(|c| &c.persister).collect(),
-	)
-}
-
-pub fn create_node_cfgs_with_persisters<'a>(
+fn create_node_cfgs_internal<'a, F>(
 	node_count: usize, chanmon_cfgs: &'a Vec<TestChanMonCfg>,
-	persisters: Vec<&'a impl test_utils::SyncPersist>,
-) -> Vec<NodeCfg<'a>> {
+	persisters: Vec<&'a impl test_utils::SyncPersist>, message_router_constructor: F,
+) -> Vec<NodeCfg<'a>>
+where
+	F: Fn(
+		Arc<NetworkGraph<&'a TestLogger>>,
+		&'a TestKeysInterface,
+	) -> test_utils::TestMessageRouter<'a>,
+{
 	let mut nodes = Vec::new();
 
 	for i in 0..node_count {
+		let cfg = &chanmon_cfgs[i];
+		let network_graph = Arc::new(NetworkGraph::new(Network::Testnet, &cfg.logger));
 		let chain_monitor = test_utils::TestChainMonitor::new(
-			Some(&chanmon_cfgs[i].chain_source),
-			&chanmon_cfgs[i].tx_broadcaster,
-			&chanmon_cfgs[i].logger,
-			&chanmon_cfgs[i].fee_estimator,
+			Some(&cfg.chain_source),
+			&cfg.tx_broadcaster,
+			&cfg.logger,
+			&cfg.fee_estimator,
 			persisters[i],
-			&chanmon_cfgs[i].keys_manager,
+			&cfg.keys_manager,
 		);
-		let network_graph = Arc::new(NetworkGraph::new(Network::Testnet, &chanmon_cfgs[i].logger));
+
 		let seed = [i as u8; 32];
 		nodes.push(NodeCfg {
-			chain_source: &chanmon_cfgs[i].chain_source,
-			logger: &chanmon_cfgs[i].logger,
-			tx_broadcaster: &chanmon_cfgs[i].tx_broadcaster,
-			fee_estimator: &chanmon_cfgs[i].fee_estimator,
+			chain_source: &cfg.chain_source,
+			logger: &cfg.logger,
+			tx_broadcaster: &cfg.tx_broadcaster,
+			fee_estimator: &cfg.fee_estimator,
 			router: test_utils::TestRouter::new(
 				Arc::clone(&network_graph),
-				&chanmon_cfgs[i].logger,
-				&chanmon_cfgs[i].scorer,
+				&cfg.logger,
+				&cfg.scorer,
 			),
-			message_router: test_utils::TestMessageRouter::new(
+			message_router: message_router_constructor(
 				Arc::clone(&network_graph),
-				&chanmon_cfgs[i].keys_manager,
+				&cfg.keys_manager,
 			),
 			chain_monitor,
-			keys_manager: &chanmon_cfgs[i].keys_manager,
+			keys_manager: &cfg.keys_manager,
 			node_seed: seed,
 			network_graph,
 			override_init_features: Rc::new(RefCell::new(None)),
@@ -4256,6 +4254,42 @@ pub fn create_node_cfgs_with_persisters<'a>(
 	nodes
 }
 
+pub fn create_node_cfgs<'a>(
+	node_count: usize, chanmon_cfgs: &'a Vec<TestChanMonCfg>,
+) -> Vec<NodeCfg<'a>> {
+	let persisters = chanmon_cfgs.iter().map(|c| &c.persister).collect();
+	create_node_cfgs_internal(
+		node_count,
+		chanmon_cfgs,
+		persisters,
+		test_utils::TestMessageRouter::new_default,
+	)
+}
+
+pub fn create_node_cfgs_with_persisters<'a>(
+	node_count: usize, chanmon_cfgs: &'a Vec<TestChanMonCfg>,
+	persisters: Vec<&'a impl test_utils::SyncPersist>,
+) -> Vec<NodeCfg<'a>> {
+	create_node_cfgs_internal(
+		node_count,
+		chanmon_cfgs,
+		persisters,
+		test_utils::TestMessageRouter::new_default,
+	)
+}
+
+pub fn create_node_cfgs_with_node_id_message_router<'a>(
+	node_count: usize, chanmon_cfgs: &'a Vec<TestChanMonCfg>,
+) -> Vec<NodeCfg<'a>> {
+	let persisters = chanmon_cfgs.iter().map(|c| &c.persister).collect();
+	create_node_cfgs_internal(
+		node_count,
+		chanmon_cfgs,
+		persisters,
+		test_utils::TestMessageRouter::new_node_id_router,
+	)
+}
+
 pub fn test_default_channel_config() -> UserConfig {
 	let mut default_config = UserConfig::default();
 	// Set cltv_expiry_delta slightly lower to keep the final CLTV values inside one byte in our

lightning/src/ln/functional_tests.rs

@@ -5175,7 +5175,7 @@ pub fn test_key_derivation_params() {
 	let router =
 		test_utils::TestRouter::new(Arc::clone(&network_graph), &chanmon_cfgs[0].logger, &scorer);
 	let message_router =
-		test_utils::TestMessageRouter::new(Arc::clone(&network_graph), &keys_manager);
+		test_utils::TestMessageRouter::new_default(Arc::clone(&network_graph), &keys_manager);
 	let node = NodeCfg {
 		chain_source: &chanmon_cfgs[0].chain_source,
 		logger: &chanmon_cfgs[0].logger,

lightning/src/util/test_utils.rs

@@ -34,7 +34,7 @@ use crate::ln::types::ChannelId;
 use crate::ln::{msgs, wire};
 use crate::offers::invoice::UnsignedBolt12Invoice;
 use crate::onion_message::messenger::{
-	DefaultMessageRouter, Destination, MessageRouter, OnionMessagePath,
+	DefaultMessageRouter, Destination, MessageRouter, NodeIdMessageRouter, OnionMessagePath,
 };
 use crate::routing::gossip::{EffectiveCapacity, NetworkGraph, NodeId, RoutingFees};
 use crate::routing::router::{
@@ -311,41 +311,69 @@ impl<'a> Drop for TestRouter<'a> {
 	}
 }
 
-pub struct TestMessageRouter<'a> {
-	inner: DefaultMessageRouter<
-		Arc<NetworkGraph<&'a TestLogger>>,
-		&'a TestLogger,
-		&'a TestKeysInterface,
-	>,
+pub enum TestMessageRouter<'a> {
+	Default(
+		DefaultMessageRouter<
+			Arc<NetworkGraph<&'a TestLogger>>,
+			&'a TestLogger,
+			&'a TestKeysInterface,
+		>,
+	),
+	NodeId(
+		NodeIdMessageRouter<
+			Arc<NetworkGraph<&'a TestLogger>>,
+			&'a TestLogger,
+			&'a TestKeysInterface,
+		>,
+	),
 }
 
 impl<'a> TestMessageRouter<'a> {
-	pub fn new(
+	pub fn new_default(
+		network_graph: Arc<NetworkGraph<&'a TestLogger>>, entropy_source: &'a TestKeysInterface,
+	) -> Self {
+		Self::Default(DefaultMessageRouter::new(network_graph, entropy_source))
+	}
+
+	pub fn new_node_id_router(
 		network_graph: Arc<NetworkGraph<&'a TestLogger>>, entropy_source: &'a TestKeysInterface,
 	) -> Self {
-		Self { inner: DefaultMessageRouter::new(network_graph, entropy_source) }
+		Self::NodeId(NodeIdMessageRouter::new(network_graph, entropy_source))
 	}
 }
 
 impl<'a> MessageRouter for TestMessageRouter<'a> {
 	fn find_path(
 		&self, sender: PublicKey, peers: Vec<PublicKey>, destination: Destination,
 	) -> Result<OnionMessagePath, ()> {
-		self.inner.find_path(sender, peers, destination)
+		match self {
+			Self::Default(inner) => inner.find_path(sender, peers, destination),
+			Self::NodeId(inner) => inner.find_path(sender, peers, destination),
+		}
 	}
 
 	fn create_blinded_paths<T: secp256k1::Signing + secp256k1::Verification>(
 		&self, recipient: PublicKey, context: MessageContext, peers: Vec<MessageForwardNode>,
 		secp_ctx: &Secp256k1<T>,
 	) -> Result<Vec<BlindedMessagePath>, ()> {
-		self.inner.create_blinded_paths(recipient, context, peers, secp_ctx)
+		match self {
+			Self::Default(inner) => inner.create_blinded_paths(recipient, context, peers, secp_ctx),
+			Self::NodeId(inner) => inner.create_blinded_paths(recipient, context, peers, secp_ctx),
+		}
 	}
 
 	fn create_compact_blinded_paths<T: secp256k1::Signing + secp256k1::Verification>(
 		&self, recipient: PublicKey, context: MessageContext, peers: Vec<MessageForwardNode>,
 		secp_ctx: &Secp256k1<T>,
 	) -> Result<Vec<BlindedMessagePath>, ()> {
-		self.inner.create_compact_blinded_paths(recipient, context, peers, secp_ctx)
+		match self {
+			Self::Default(inner) => {
+				inner.create_compact_blinded_paths(recipient, context, peers, secp_ctx)
+			},
+			Self::NodeId(inner) => {
+				inner.create_compact_blinded_paths(recipient, context, peers, secp_ctx)
+			},
+		}
 	}
 }