From 264f181c82effdc5b160997538186c9c1542e86f Mon Sep 17 00:00:00 2001
From: jurvis <hello@jurvis.co>
Date: Thu, 14 Jul 2022 19:40:16 -0700
Subject: [PATCH] Add utils to handle HTLC handling failure reason

We add `HTLCHandlingFailedConditions` to express the failure parameters, that will be enforced by a new macro, `expect_pending_htlcs_forwardable_conditions`.
---
 lightning/src/ln/functional_test_utils.rs | 91 +++++++++++++++++++++--
 1 file changed, 85 insertions(+), 6 deletions(-)

diff --git a/lightning/src/ln/functional_test_utils.rs b/lightning/src/ln/functional_test_utils.rs
index fe78395e2f9..c4f2480eb18 100644
--- a/lightning/src/ln/functional_test_utils.rs
+++ b/lightning/src/ln/functional_test_utils.rs
@@ -24,7 +24,7 @@ use util::enforcing_trait_impls::EnforcingSigner;
 use util::scid_utils;
 use util::test_utils;
 use util::test_utils::{panicking, TestChainMonitor};
-use util::events::{Event, MessageSendEvent, MessageSendEventsProvider, PaymentPurpose};
+use util::events::{Event, HTLCDestination, MessageSendEvent, MessageSendEventsProvider, PaymentPurpose};
 use util::errors::APIError;
 use util::config::UserConfig;
 use util::ser::{ReadableArgs, Writeable};
@@ -1254,24 +1254,101 @@ macro_rules! get_route_and_payment_hash {
 	}}
 }
 
+pub struct HTLCHandlingFailedConditions {
+	pub expected_pending_htlcs_forwardable: bool,
+	pub expected_htlc_processing_failed: Option<u32>,
+	pub expected_destination: Option<HTLCDestination>,
+}
+
+impl HTLCHandlingFailedConditions {
+	pub fn new() -> Self {
+		Self {
+			expected_pending_htlcs_forwardable: false,
+			expected_htlc_processing_failed: None,
+			expected_destination: None,
+		}
+	}
+
+	pub fn htlc_processing_failed_with_count(mut self, count: u32) -> Self {
+		self.expected_htlc_processing_failed = Some(count);
+		self
+	}
+
+	pub fn htlc_processing_failed_with_reason(mut self, reason: HTLCDestination) -> Self {
+		self.expected_htlc_processing_failed = Some(1);
+		self.expected_destination = Some(reason);
+		self
+	}
+
+	pub fn htlc_processing_failed_with_count_and_reason(mut self, count: u32, reason: HTLCDestination) -> Self {
+		self.expected_htlc_processing_failed = Some(count);
+		self.expected_destination = Some(reason);
+		self
+	}
+}
+
 #[macro_export]
-/// Clears (and ignores) a PendingHTLCsForwardable event
-macro_rules! expect_pending_htlcs_forwardable_ignore {
-	($node: expr) => {{
+macro_rules! expect_pending_htlcs_forwardable_conditions {
+	($node: expr, $conditions: expr) => {{
 		let events = $node.node.get_and_clear_pending_events();
-		assert_eq!(events.len(), 1);
 		match events[0] {
 			$crate::util::events::Event::PendingHTLCsForwardable { .. } => { },
 			_ => panic!("Unexpected event"),
 		};
+
+		if let Some(count) = $conditions.expected_htlc_processing_failed {
+			assert_eq!(events.len() as u32, count + 1u32);
+		} else {
+			assert_eq!(events.len(), 1);
+		}
+
+		if let Some(reason) = $conditions.expected_destination {
+			for event in events {
+				match event {
+					$crate::util::events::Event::PendingHTLCsForwardable { .. } => { },
+					$crate::util::events::Event::HTLCHandlingFailed { ref failed_next_destination, .. } => {
+						assert_eq!(reason, failed_next_destination.clone());
+					},
+					_ => panic!("Unexpected event"),
+				}
+			}
+		}
 	}}
 }
 
+#[macro_export]
+/// Clears (and ignores) a PendingHTLCsForwardable event
+macro_rules! expect_pending_htlcs_forwardable_ignore {
+	($node: expr) => {{
+		expect_pending_htlcs_forwardable_conditions!($node, $crate::ln::functional_test_utils::HTLCHandlingFailedConditions::new());
+	}};
+}
+
+#[macro_export]
+/// Clears (and ignores) PendingHTLCsForwardable and HTLCHandlingFailed events
+macro_rules! expect_pending_htlcs_forwardable_and_htlc_handling_failed_ignore {
+	($node: expr, $conditions: expr) => {{
+		expect_pending_htlcs_forwardable_conditions!($node, $conditions);
+	}};
+}
+
 #[macro_export]
 /// Handles a PendingHTLCsForwardable event
 macro_rules! expect_pending_htlcs_forwardable {
 	($node: expr) => {{
-		$crate::expect_pending_htlcs_forwardable_ignore!($node);
+		expect_pending_htlcs_forwardable_ignore!($node);
+		$node.node.process_pending_htlc_forwards();
+
+		// Ensure process_pending_htlc_forwards is idempotent.
+		$node.node.process_pending_htlc_forwards();
+	}};
+}
+
+#[macro_export]
+/// Handles a PendingHTLCsForwardable and HTLCHandlingFailed event
+macro_rules! expect_pending_htlcs_forwardable_and_htlc_handling_failed {
+	($node: expr, $conditions: expr) => {{
+		expect_pending_htlcs_forwardable_and_htlc_handling_failed_ignore!($node, $conditions);
 		$node.node.process_pending_htlc_forwards();
 
 		// Ensure process_pending_htlc_forwards is idempotent.
@@ -1282,6 +1359,8 @@ macro_rules! expect_pending_htlcs_forwardable {
 #[cfg(test)]
 macro_rules! expect_pending_htlcs_forwardable_from_events {
 	($node: expr, $events: expr, $ignore: expr) => {{
+		// We need to clear pending events since there may possibly be `PaymentForwardingFailed` events here
+		$node.node.get_and_clear_pending_events();
 		assert_eq!($events.len(), 1);
 		match $events[0] {
 			Event::PendingHTLCsForwardable { .. } => { },