Add a test for asserting existing inflight amounts

Author: jurvis

lightning-invoice/src/payment.rs

@@ -801,7 +801,7 @@ mod tests {
 	use lightning::ln::features::{ChannelFeatures, NodeFeatures, InitFeatures};
 	use lightning::ln::functional_test_utils::*;
 	use lightning::ln::msgs::{ChannelMessageHandler, ErrorAction, LightningError};
-	use lightning::routing::gossip::NodeId;
+	use lightning::routing::gossip::{EffectiveCapacity, NodeId};
 	use lightning::routing::router::{PaymentParameters, Route, RouteHop};
 	use lightning::routing::scoring::ChannelUsage;
 	use lightning::util::test_utils::TestLogger;
@@ -1550,7 +1550,7 @@ mod tests {
 	}
 
 	#[test]
-	fn inflight_data_trivial_test() {
+	fn inflight_map_data_trivial_test() {
 		let event_handled = core::cell::RefCell::new(false);
 		let event_handler = |_: &_| { *event_handled.borrow_mut() = true; };
 
@@ -1590,7 +1590,57 @@ mod tests {
 		assert_eq!(inflight_map.get(&(2, true)), None);
 
 		// Second path should still be inflight
-		assert_eq!(inflight_map.get(&(1, true)).unwrap().clone(), 64);
+		assert_eq!(inflight_map.get(&(1, true)).unwrap().clone(), 64)
+	}
+
+	#[test]
+	fn inflight_usage_trivial_test() {
+		// First, let's just send a payment through, but only make sure one of the path completes
+		let event_handled = core::cell::RefCell::new(false);
+		let event_handler = |_: &_| { *event_handled.borrow_mut() = true; };
+
+		let payment_preimage = PaymentPreimage([1; 32]);
+		let payment_invoice = invoice(payment_preimage);
+		let payment_hash = Some(PaymentHash(payment_invoice.payment_hash().clone().into_inner()));
+		let final_value_msat = payment_invoice.amount_milli_satoshis().unwrap();
+
+		let payer = TestPayer::new()
+			.expect_send(Amount::ForInvoice(128))
+			.expect_send(Amount::ForInvoice(128));
+		let final_value_msat = payment_invoice.amount_milli_satoshis().unwrap();
+		let route = TestRouter::route_for_value(final_value_msat);
+		let router = TestRouter {};
+		let scorer = RefCell::new(TestScorer::new()
+			// 1st invoice, 1st path
+	  		.expect_usage(ChannelUsage { amount_msat: 10, inflight_htlc_msat: 0, effective_capacity: EffectiveCapacity::Unknown } )
+			.expect_usage(ChannelUsage { amount_msat: 20, inflight_htlc_msat: 0, effective_capacity: EffectiveCapacity::Unknown } )
+			.expect_usage(ChannelUsage { amount_msat: 64, inflight_htlc_msat: 0, effective_capacity: EffectiveCapacity::Unknown } )
+			// 1st invoice, 2nd path
+			.expect_usage(ChannelUsage { amount_msat: 64, inflight_htlc_msat: 0, effective_capacity: EffectiveCapacity::Unknown } )
+			// 2nd invoice, 1st path
+			.expect_usage(ChannelUsage { amount_msat: 10, inflight_htlc_msat: 0, effective_capacity: EffectiveCapacity::Unknown } )
+		    .expect_usage(ChannelUsage { amount_msat: 20, inflight_htlc_msat: 0, effective_capacity: EffectiveCapacity::Unknown } )
+		    .expect_usage(ChannelUsage { amount_msat: 64, inflight_htlc_msat: 0, effective_capacity: EffectiveCapacity::Unknown } )
+			// 2nd invoice, 2nd path
+			.expect_usage(ChannelUsage { amount_msat: 128, inflight_htlc_msat: 0, effective_capacity: EffectiveCapacity::Unknown } )
+		);
+		let logger = TestLogger::new();
+		let invoice_payer =
+			InvoicePayer::new(&payer, router, &scorer, &logger, event_handler, Retry::Attempts(0));
+
+		// Succeed 1st path, leave 2nd path inflight
+		let payment_id = invoice_payer.pay_invoice(&payment_invoice).unwrap();
+		invoice_payer.handle_event(&Event::PaymentPathSuccessful {
+			payment_id, payment_hash, path: route.paths[0].clone()
+		});
+
+		// Let's pay a second invoice that will be using the same path. This should trigger the
+		// assertions that expect the last 4 ChannelUsage values above where TestScorer is initialized.
+		// Particularly, the 2nd path of the 1st payment, since it is not yet complete, should still
+		// have 64 msats inflight. Making the total inflight 128.
+		let payment_preimage_2 = PaymentPreimage([2; 32]);
+		let payment_invoice_2 = invoice(payment_preimage_2);
+		invoice_payer.pay_invoice(&payment_invoice_2).unwrap();
 	}
 
 	struct TestRouter;
@@ -1658,9 +1708,22 @@ mod tests {
 
 	impl Router for TestRouter {
 		fn find_route<S: Score>(
-			&self, _payer: &PublicKey, route_params: &RouteParameters, _payment_hash: &PaymentHash,
-			_first_hops: Option<&[&ChannelDetails]>, _scorer: &S
+			&self, payer: &PublicKey, route_params: &RouteParameters, _payment_hash: &PaymentHash,
+			_first_hops: Option<&[&ChannelDetails]>, scorer: &S
 		) -> Result<Route, LightningError> {
+			// Simulate calling the Scorer just as you would in find_route
+			let route = Self::route_for_value(route_params.final_value_msat);
+			for path in route.paths {
+				for hop in path {
+					let usage = ChannelUsage {
+						amount_msat: hop.fee_msat,
+						inflight_htlc_msat: 0,
+						effective_capacity: EffectiveCapacity::Unknown,
+					};
+					scorer.channel_penalty_msat(hop.short_channel_id, &NodeId::from_pubkey(payer), &NodeId::from_pubkey(&hop.pubkey), usage);
+				}
+			}
+
 			Ok(Route {
 				payment_params: Some(route_params.payment_params.clone()), ..Self::route_for_value(route_params.final_value_msat)
 			})
@@ -1680,6 +1743,7 @@ mod tests {
 
 	struct TestScorer {
 		expectations: Option<VecDeque<TestResult>>,
+		scorer_expectations: RefCell<VecDeque<ChannelUsage>>,
 	}
 
 	#[derive(Debug)]
@@ -1694,13 +1758,19 @@ mod tests {
 		fn new() -> Self {
 			Self {
 				expectations: None,
+				scorer_expectations: RefCell::new(VecDeque::new()),
 			}
 		}
 
 		fn expect(mut self, expectation: TestResult) -> Self {
 			self.expectations.get_or_insert_with(|| VecDeque::new()).push_back(expectation);
 			self
 		}
+
+		fn expect_usage(self, expectation: ChannelUsage) -> Self {
+			self.scorer_expectations.borrow_mut().push_back(expectation);
+			self
+		}
 	}
 
 	#[cfg(c_bindings)]
@@ -1710,8 +1780,16 @@ mod tests {
 
 	impl Score for TestScorer {
 		fn channel_penalty_msat(
-			&self, _short_channel_id: u64, _source: &NodeId, _target: &NodeId, _usage: ChannelUsage
-		) -> u64 { 0 }
+			&self, _short_channel_id: u64, _source: &NodeId, _target: &NodeId, usage: ChannelUsage
+		) -> u64 {
+			match self.scorer_expectations.borrow_mut().pop_front() {
+				Some(expectation) => {
+					assert_eq!(expectation.amount_msat, usage.amount_msat);
+				},
+				None => {},
+			}
+			0
+		}
 
 		fn payment_path_failed(&mut self, actual_path: &[&RouteHop], actual_short_channel_id: u64) {
 			if let Some(expectations) = &mut self.expectations {