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

@@ -17077,7 +17077,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

@@ -38,7 +38,7 @@ use crate::util::logger::Logger;
 use crate::util::scid_utils;
 use crate::util::test_channel_signer::TestChannelSigner;
 use crate::util::test_channel_signer::SignerOp;
-use crate::util::test_utils;
+use crate::util::test_utils::{self, TestLogger};
 use crate::util::test_utils::{TestChainMonitor, TestScorer, TestKeysInterface};
 use crate::util::ser::{ReadableArgs, Writeable};
 
@@ -728,7 +728,7 @@ impl<'a, 'b, 'c> Drop for Node<'a, 'b, 'c> {
 					signer_provider: self.keys_manager,
 					fee_estimator: &test_utils::TestFeeEstimator::new(253),
 					router: &test_utils::TestRouter::new(Arc::clone(&network_graph), &self.logger, &scorer),
-					message_router: &test_utils::TestMessageRouter::new(network_graph, self.keys_manager),
+					message_router: &test_utils::TestMessageRouter::new_default(network_graph, self.keys_manager),
 					chain_monitor: self.chain_monitor,
 					tx_broadcaster: &broadcaster,
 					logger: &self.logger,
@@ -3345,26 +3345,39 @@ pub fn create_chanmon_cfgs_with_keys(node_count: usize, predefined_keys_ids: Opt
 	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>(node_count: usize, chanmon_cfgs: &'a Vec<TestChanMonCfg>, persisters: Vec<&'a impl test_utils::SyncPersist>) -> Vec<NodeCfg<'a>> {
+fn create_node_cfgs_internal<'a, F>(
+	node_count: usize,
+	chanmon_cfgs: &'a Vec<TestChanMonCfg>,
+	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 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, persisters[i], &chanmon_cfgs[i].keys_manager);
-		let network_graph = Arc::new(NetworkGraph::new(Network::Testnet, &chanmon_cfgs[i].logger));
+		let cfg = &chanmon_cfgs[i];
+		let network_graph = Arc::new(NetworkGraph::new(Network::Testnet, &cfg.logger));
+		let chain_monitor = test_utils::TestChainMonitor::new(
+			Some(&cfg.chain_source),
+			&cfg.tx_broadcaster,
+			&cfg.logger,
+			&cfg.fee_estimator,
+			persisters[i],
+			&cfg.keys_manager,
+		);
+
 		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,
-			router: test_utils::TestRouter::new(Arc::clone(&network_graph), &chanmon_cfgs[i].logger, &chanmon_cfgs[i].scorer),
-			message_router: test_utils::TestMessageRouter::new(Arc::clone(&network_graph), &chanmon_cfgs[i].keys_manager),
+			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), &cfg.logger, &cfg.scorer),
+			message_router: message_router_constructor(Arc::clone(&network_graph), &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)),
@@ -3374,6 +3387,33 @@ pub fn create_node_cfgs_with_persisters<'a>(node_count: usize, chanmon_cfgs: &'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, |graph, keys_manager| {
+		test_utils::TestMessageRouter::new_default(graph, keys_manager)
+	})
+}
+
+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, |graph, keys_manager| {
+		test_utils::TestMessageRouter::new_default(graph, keys_manager)
+	})
+}
+
+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, |graph, keys_manager| {
+		test_utils::TestMessageRouter::new_node_id_router(graph, keys_manager)
+	})
+}
+
 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,74 @@ 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 {
+		inner: DefaultMessageRouter<
+			Arc<NetworkGraph<&'a TestLogger>>,
+			&'a TestLogger,
+			&'a TestKeysInterface,
+		>,
+	},
+
+	NodeId {
+		inner: 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 { inner: DefaultMessageRouter::new(network_graph, entropy_source) }
+		Self::Default { inner: DefaultMessageRouter::new(network_graph, entropy_source) }
+	}
+
+	pub fn new_node_id_router(
+		network_graph: Arc<NetworkGraph<&'a TestLogger>>, entropy_source: &'a TestKeysInterface,
+	) -> Self {
+		Self::NodeId { inner: 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)
+			},
+		}
 	}
 }