add helper functions for probes

Author: alexanderwiederin

lightning/src/ln/functional_test_utils.rs

@@ -2078,6 +2078,25 @@ pub fn expect_channel_ready_event<'a, 'b, 'c, 'd>(node: &'a Node<'b, 'c, 'd>, ex
 	}
 }
 
+#[cfg(any(test, feature = "_test_utils"))]
+pub fn expect_probe_successful_events(node: &Node, probe_results: Vec<(PaymentHash, PaymentId)>) {
+	let mut events = node.node.get_and_clear_pending_events();
+
+	assert_eq!(events.len(), probe_results.len());
+	for (idx, event) in events.drain(..).enumerate() {
+		let (payment_hash, payment_id) = probe_results[idx];
+		match event {
+			Event::ProbeSuccessful { payment_id: ev_pid, payment_hash: ev_ph, ..} => {
+				assert_eq!(payment_id, ev_pid);
+				assert_eq!(payment_hash, ev_ph);
+			},
+			_ => panic!(),
+		}
+	};
+
+	assert!(!node.node.has_pending_payments());
+}
+
 pub struct PaymentFailedConditions<'a> {
 	pub(crate) expected_htlc_error_data: Option<(u16, &'a [u8])>,
 	pub(crate) expected_blamed_scid: Option<u64>,
@@ -2215,6 +2234,67 @@ pub fn send_along_route_with_secret<'a, 'b, 'c>(origin_node: &Node<'a, 'b, 'c>,
 	payment_id
 }
 
+pub fn fail_payment_along_path<'a, 'b, 'c>(expected_path: &[&Node<'a, 'b, 'c>]) {
+	let origin_node_id = expected_path[0].node.get_our_node_id();
+
+	// iterate from the receiving node to the origin node and handle update fail htlc.
+	for (&node, &prev_node) in expected_path.iter().rev().zip(expected_path.iter().rev().skip(1)) {
+		let updates = get_htlc_update_msgs!(node, prev_node.node.get_our_node_id());
+		prev_node.node.handle_update_fail_htlc(&node.node.get_our_node_id(), &updates.update_fail_htlcs[0]);
+		check_added_monitors!(prev_node, 0);
+
+		if origin_node_id == prev_node.node.get_our_node_id() {
+			commitment_signed_dance!(prev_node, node, updates.commitment_signed, false);
+		} else {
+			commitment_signed_dance!(prev_node, node, updates.commitment_signed, true);
+		}
+	}
+}
+
+pub fn pass_probe_along_route<'a, 'b, 'c>(origin_node: &Node<'a, 'b, 'c>, expected_route: &[&[&Node<'a, 'b, 'c>]]) {
+	let mut events = origin_node.node.get_and_clear_pending_msg_events();
+	assert_eq!(events.len(), expected_route.len());
+
+	check_added_monitors!(origin_node, expected_route.len());
+	for path in expected_route.iter() {
+		let ev = remove_first_msg_event_to_node(&path[0].node.get_our_node_id(), &mut events);
+
+		pass_probe_along_path(origin_node, path, ev);
+		let nodes_to_fail_payment: Vec<_> = vec![origin_node].into_iter().chain(path.iter().cloned()).collect();
+
+		fail_payment_along_path(nodes_to_fail_payment.as_slice());
+	}
+}
+
+
+pub fn pass_probe_along_path<'a, 'b, 'c>(origin_node: &Node<'a, 'b, 'c>, expected_path: &[&Node<'a, 'b, 'c>], ev: MessageSendEvent) {
+	let mut payment_event = SendEvent::from_event(ev);
+	let mut prev_node = origin_node;
+
+	for (idx, &node) in expected_path.iter().enumerate() {
+		assert_eq!(node.node.get_our_node_id(), payment_event.node_id);
+
+		node.node.handle_update_add_htlc(&prev_node.node.get_our_node_id(), &payment_event.msgs[0]);
+		check_added_monitors!(node, 0);
+
+		if idx != expected_path.len() - 1 {
+			commitment_signed_dance!(node, prev_node, payment_event.commitment_msg, false);
+			expect_pending_htlcs_forwardable!(node);
+
+			let mut events = node.node.get_and_clear_pending_msg_events();
+			assert_eq!(events.len(), 1);
+			payment_event = SendEvent::from_event(events.remove(0));
+			assert_eq!(payment_event.msgs.len(), 1);
+			check_added_monitors!(node, 1);
+		} else {
+			commitment_signed_dance!(node, prev_node, payment_event.commitment_msg, true, true);
+		}
+
+
+		prev_node = node;
+	}
+}
+
 pub fn do_pass_along_path<'a, 'b, 'c>(origin_node: &Node<'a, 'b, 'c>, expected_path: &[&Node<'a, 'b, 'c>], recv_value: u64, our_payment_hash: PaymentHash, our_payment_secret: Option<PaymentSecret>, ev: MessageSendEvent, payment_claimable_expected: bool, clear_recipient_events: bool, expected_preimage: Option<PaymentPreimage>) -> Option<Event> {
 	let mut payment_event = SendEvent::from_event(ev);
 	let mut prev_node = origin_node;