Expand test to include mpp

Author: sourabhmarathe

lightning/src/routing/network_graph.rs

@@ -1753,16 +1753,30 @@ mod tests {
 
 	#[test]
 	fn network_graph_score_payment() {
-		// Set up a network that consists of just a channel between node_1 and node_2
+		// Set up a network that consists of the following channels:
+		// node_1 -> node_2
+		// node_1 -> node_3
+		// node_2 -> node_4
+		// node_3 -> node_4
+		// We can use this to simulate a MPP from node_1 to node_4 via node_2 and node_3.
 		let (secp_ctx, net_graph_msg_handler) = create_net_graph_msg_handler();
-		let node_1_privkey = &SecretKey::from_slice(&[42; 32]).unwrap();
-		let node_2_privkey = &SecretKey::from_slice(&[41; 32]).unwrap();
+		let node_0_privkey = &SecretKey::from_slice(&[42; 32]).unwrap();
+		let node_1_privkey = &SecretKey::from_slice(&[41; 32]).unwrap();
+		let node_2_privkey = &SecretKey::from_slice(&[40; 32]).unwrap();
+		let node_3_privkey = &SecretKey::from_slice(&[39; 32]).unwrap();
+		let node_privkeys = vec![node_0_privkey, node_1_privkey, node_2_privkey, node_3_privkey];
+		let node_id_0 = PublicKey::from_secret_key(&secp_ctx, node_0_privkey);
 		let node_id_1 = PublicKey::from_secret_key(&secp_ctx, node_1_privkey);
 		let node_id_2 = PublicKey::from_secret_key(&secp_ctx, node_2_privkey);
-		let node_1_btckey = &SecretKey::from_slice(&[40; 32]).unwrap();
-		let node_2_btckey = &SecretKey::from_slice(&[39; 32]).unwrap();
-
-		let short_channel_id = 0;
+		let node_id_3 = PublicKey::from_secret_key(&secp_ctx, node_3_privkey);
+		let node_ids = vec![node_id_0, node_id_1, node_id_2, node_id_3];
+		let node_0_btckey = &SecretKey::from_slice(&[38; 32]).unwrap();
+		let node_1_btckey = &SecretKey::from_slice(&[37; 32]).unwrap();
+		let node_2_btckey = &SecretKey::from_slice(&[36; 32]).unwrap();
+		let node_3_btckey = &SecretKey::from_slice(&[35; 32]).unwrap();
+		let node_btckeys = vec![node_0_btckey, node_1_btckey, node_2_btckey, node_3_btckey];
+
+		let mut short_channel_id = 0;
 		let chain_hash = genesis_block(Network::Testnet).header.bitcoin_hash();
 
 		{
@@ -1772,84 +1786,103 @@ mod tests {
 		}
 
 		{
-			// Announce a channel we will update
-			let unsigned_announcement = UnsignedChannelAnnouncement {
-				features: ChannelFeatures::empty(),
-				chain_hash,
-				short_channel_id,
-				node_id_1,
-				node_id_2,
-				bitcoin_key_1: PublicKey::from_secret_key(&secp_ctx, node_1_btckey),
-				bitcoin_key_2: PublicKey::from_secret_key(&secp_ctx, node_2_btckey),
-				excess_data: Vec::new(),
-			};
+			// Create channels between each node
+			for i in 0..=3 {
+				for j in 0..=3 {
+					if i >= j {
+						continue;
+					} else {
+						let node_i = node_ids[i];
+						let node_j = node_ids[j];
+						let unsigned_announcement = UnsignedChannelAnnouncement {
+							features: ChannelFeatures::empty(),
+							chain_hash,
+							short_channel_id,
+							node_id_1: node_i,
+							node_id_2: node_j,
+							bitcoin_key_1: PublicKey::from_secret_key(&secp_ctx, node_btckeys[i]),
+							bitcoin_key_2: PublicKey::from_secret_key(&secp_ctx, node_btckeys[j]),
+							excess_data: Vec::new(),
+						};
 
-			let msghash = hash_to_message!(&Sha256dHash::hash(&unsigned_announcement.encode()[..])[..]);
-			let valid_channel_announcement = ChannelAnnouncement {
-				node_signature_1: secp_ctx.sign(&msghash, node_1_privkey),
-				node_signature_2: secp_ctx.sign(&msghash, node_2_privkey),
-				bitcoin_signature_1: secp_ctx.sign(&msghash, node_1_btckey),
-				bitcoin_signature_2: secp_ctx.sign(&msghash, node_2_btckey),
-				contents: unsigned_announcement.clone(),
-			};
-			match net_graph_msg_handler.handle_channel_announcement(&valid_channel_announcement) {
-				Ok(_) => (),
-				Err(_) => panic!()
-			};
+						let msghash = hash_to_message!(&Sha256dHash::hash(&unsigned_announcement.encode()[..])[..]);
+						let valid_channel_announcement = ChannelAnnouncement {
+							node_signature_1: secp_ctx.sign(&msghash, node_privkeys[i]),
+							node_signature_2: secp_ctx.sign(&msghash, node_privkeys[j]),
+							bitcoin_signature_1: secp_ctx.sign(&msghash, node_btckeys[i]),
+							bitcoin_signature_2: secp_ctx.sign(&msghash, node_btckeys[j]),
+							contents: unsigned_announcement.clone(),
+						};
+						match net_graph_msg_handler.handle_channel_announcement(&valid_channel_announcement) {
+							Ok(_) => (),
+							Err(_) => panic!()
+						};
+						short_channel_id += 1;
+					}
+				}
+			}
+			// There is no nodes in the table at the beginning.
+			let network = net_graph_msg_handler.network_graph.read().unwrap();
+			assert_eq!(network.get_nodes().len(), 4);
+			assert_eq!(network.get_channels().len(), 6);
 		}
 
 		{
-			// At the start, the node_1 <-> node_2 channel's DirectionalChannelInfos should be 0.
+			// Create the DirectionalChannelInfo, RoutingFee and ChannelScore objects for each
+			// channel
 			let mut network = net_graph_msg_handler.network_graph.write().unwrap();
-			let chan_id = network.get_nodes().get(&node_id_1).unwrap().channels[0];
-			let channel_info = network.channels.get_mut(&chan_id).unwrap();
-			// Assign a DirectionalChannelInfo and ChannelScore object to alice_to_bob of
-			// channel_info
-			let chan_score_1 = ChannelScore {
-				payment_sent_score: 0,
-				payment_failed_score: 0,
-			};
-			let chan_score_2 = ChannelScore {
-				payment_sent_score: 0,
-				payment_failed_score: 0,
-			};
+			for channel_id in 0..short_channel_id {
+				//let chan_id = network.get_nodes().get(&node_ids[0]).unwrap().channels[0];
+				let channel_info = network.channels.get_mut(&channel_id).unwrap();
+				// Assign a DirectionalChannelInfo and ChannelScore object to alice_to_bob of
+				// channel_info
+				let chan_score_alice_to_bob = ChannelScore {
+					payment_sent_score: 0,
+					payment_failed_score: 0,
+				};
+				let chan_score_bob_to_alice = ChannelScore {
+					payment_sent_score: 0,
+					payment_failed_score: 0,
+				};
 
-			let dummy_routing_fee = RoutingFees {
-				base_msat: 0,
-				proportional_millionths: 0,
-			};
-			let dir_chan_info_1 = DirectionalChannelInfo {
-				last_update: 0,
-				enabled: true,
-				cltv_expiry_delta: 0,
-				htlc_minimum_msat: 0,
-				fees: dummy_routing_fee,
-				last_update_message: None,
-				channel_score: chan_score_1,
-			};
+				let dummy_routing_fee = RoutingFees {
+					base_msat: 0,
+					proportional_millionths: 0,
+				};
+				let dir_chan_info_alice_to_bob = DirectionalChannelInfo {
+					last_update: 0,
+					enabled: true,
+					cltv_expiry_delta: 0,
+					htlc_minimum_msat: 0,
+					fees: dummy_routing_fee,
+					last_update_message: None,
+					channel_score: chan_score_alice_to_bob,
+				};
 
-			let dir_chan_info_2 = DirectionalChannelInfo {
-				last_update: 0,
-				enabled: true,
-				cltv_expiry_delta: 0,
-				htlc_minimum_msat: 0,
-				fees: dummy_routing_fee,
-				last_update_message: None,
-				channel_score: chan_score_2,
-			};
+				let dir_chan_info_bob_to_alice = DirectionalChannelInfo {
+					last_update: 0,
+					enabled: true,
+					cltv_expiry_delta: 0,
+					htlc_minimum_msat: 0,
+					fees: dummy_routing_fee,
+					last_update_message: None,
+					channel_score: chan_score_bob_to_alice,
+				};
 
-			channel_info.alice_to_bob = Some(dir_chan_info_1);
-			channel_info.bob_to_alice = Some(dir_chan_info_2);
-			let dir_alice_to_bob = channel_info.alice_to_bob.as_mut().unwrap();
-			assert_eq!(dir_alice_to_bob.channel_score.payment_sent_score, 0);
-			assert_eq!(dir_alice_to_bob.channel_score.payment_failed_score, 0);
+				channel_info.alice_to_bob = Some(dir_chan_info_alice_to_bob);
+				channel_info.bob_to_alice = Some(dir_chan_info_bob_to_alice);
+				let dir_alice_to_bob = channel_info.alice_to_bob.as_mut().unwrap();
+				assert_eq!(dir_alice_to_bob.channel_score.payment_sent_score, 0);
+				assert_eq!(dir_alice_to_bob.channel_score.payment_failed_score, 0);
+			}
 		}
 
+		// From here on out we simulate the network learning of payments.
 		{
-			// Assume a payment is made, and node 1 is learning of a PaymentSent event. It now calls
-			// its score_payment_sent_for_route on itself.
+			// Assume a PaymentSent event is picked up and it consists of a simple payment from
+			// node_0 to node_1
 			let route_hop = RouteHop {
-				pubkey: node_id_1,
+				pubkey: node_ids[0],
 				node_features: NodeFeatures::empty(),
 				short_channel_id: 0,
 				channel_features: ChannelFeatures::empty(),
@@ -1862,7 +1895,7 @@ mod tests {
 			let sent_res = network.score_payment_sent_for_route(sent_route);
 			assert!(sent_res.unwrap() == true);
 			// Check that score_payment_sent_for_route incremented the appropriate ChannelScore
-			let chan_id = network.get_nodes().get(&node_id_1).unwrap().channels[0];
+			let chan_id = network.get_nodes().get(&node_ids[0]).unwrap().channels[0];
 			let channel_info = network.channels.get_mut(&chan_id).unwrap();
 			let dir_alice_to_bob = &channel_info.alice_to_bob.as_ref();
 			assert_eq!(dir_alice_to_bob.unwrap().channel_score.payment_sent_score, 1);
@@ -1871,7 +1904,7 @@ mod tests {
 
 		{
 			let route_hop = RouteHop {
-								pubkey: node_id_1,
+								pubkey: node_ids[0],
 								node_features: NodeFeatures::empty(),
 								short_channel_id: 0,
 								channel_features: ChannelFeatures::empty(),
@@ -1881,15 +1914,15 @@ mod tests {
 			// Assume a payment fails due to node_1 (identified by its PublicKey), verify that its
 			// directional channel's score (node_1 -> node_2) has its PaymentFailed score
 			// incremented.
-			let faultive_nodes: Vec<PublicKey> = vec![node_id_1];
+			let faultive_nodes: Vec<PublicKey> = vec![node_ids[0]];
 			let attempted_path: Vec<RouteHop> = vec![route_hop];
 			let failed_route: Vec<Vec<RouteHop>> = vec![attempted_path];
 			let mut network = net_graph_msg_handler.network_graph.write().unwrap();
 			let failed_res = network.score_payment_failed_for_route(failed_route, faultive_nodes);
 			assert!(failed_res.unwrap() == true);
-			let chan_id = network.get_nodes().get(&node_id_1).unwrap().channels[0];
+			let chan_id = network.get_nodes().get(&node_ids[0]).unwrap().channels[0];
 			let channel_info = network.channels.get_mut(&chan_id).unwrap();
-			assert!(channel_info.node_two == node_id_2);
+			assert!(channel_info.node_two == node_ids[1]);
 			let dir_alice_to_bob = &channel_info.alice_to_bob.as_ref();
 			assert_eq!(dir_alice_to_bob.unwrap().channel_score.payment_sent_score, 1);
 			assert_eq!(dir_alice_to_bob.unwrap().channel_score.payment_failed_score, 1);