Skip to content

Commit fad7d49

Browse files
committed
Introduce custom_data test framework for message::ReceiveTlvs
1 parent 96f6719 commit fad7d49

File tree

1 file changed

+17
-3
lines changed

1 file changed

+17
-3
lines changed

lightning/src/onion_message/functional_tests.rs

Lines changed: 17 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -157,6 +157,7 @@ struct TestCustomMessageHandler {
157157
struct OnHandleCustomMessage {
158158
expect: TestCustomMessage,
159159
include_reply_path: bool,
160+
custom_data: Option<Vec<u8>>,
160161
}
161162

162163
impl TestCustomMessageHandler {
@@ -165,13 +166,25 @@ impl TestCustomMessageHandler {
165166
}
166167

167168
fn expect_message(&self, message: TestCustomMessage) {
168-
let expectation = OnHandleCustomMessage { expect: message, include_reply_path: false };
169+
let expectation =
170+
OnHandleCustomMessage { expect: message, include_reply_path: false, custom_data: None };
169171
self.expectations.lock().unwrap().push_back(expectation);
170172
}
171173

172174
fn expect_message_and_response(&self, message: TestCustomMessage) {
173-
let expectation = OnHandleCustomMessage { expect: message, include_reply_path: true };
174-
self.expectations.lock().unwrap().push_back(expectation);
175+
self.expectations.lock().unwrap().push_back(OnHandleCustomMessage {
176+
expect: message,
177+
include_reply_path: true,
178+
custom_data: None,
179+
});
180+
}
181+
182+
fn expect_message_with_custom_data(&self, message: TestCustomMessage, custom_data: Vec<u8>) {
183+
self.expectations.lock().unwrap().push_back(OnHandleCustomMessage {
184+
expect: message,
185+
include_reply_path: false,
186+
custom_data: Some(custom_data),
187+
});
175188
}
176189

177190
fn get_next_expectation(&self) -> OnHandleCustomMessage {
@@ -196,6 +209,7 @@ impl CustomOnionMessageHandler for TestCustomMessageHandler {
196209
) -> Option<(Self::CustomMessage, ResponseInstruction)> {
197210
let expectation = self.get_next_expectation();
198211
assert_eq!(msg, expectation.expect);
212+
assert_eq!(custom_data, expectation.custom_data);
199213

200214
let response = match msg {
201215
TestCustomMessage::Ping => TestCustomMessage::Pong,

0 commit comments

Comments
 (0)