@@ -157,6 +157,7 @@ struct TestCustomMessageHandler {
157
157
struct OnHandleCustomMessage {
158
158
expect : TestCustomMessage ,
159
159
include_reply_path : bool ,
160
+ custom_data : Option < Vec < u8 > > ,
160
161
}
161
162
162
163
impl TestCustomMessageHandler {
@@ -165,13 +166,25 @@ impl TestCustomMessageHandler {
165
166
}
166
167
167
168
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 } ;
169
171
self . expectations . lock ( ) . unwrap ( ) . push_back ( expectation) ;
170
172
}
171
173
172
174
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
+ } ) ;
175
188
}
176
189
177
190
fn get_next_expectation ( & self ) -> OnHandleCustomMessage {
@@ -196,6 +209,7 @@ impl CustomOnionMessageHandler for TestCustomMessageHandler {
196
209
) -> Option < ( Self :: CustomMessage , ResponseInstruction ) > {
197
210
let expectation = self . get_next_expectation ( ) ;
198
211
assert_eq ! ( msg, expectation. expect) ;
212
+ assert_eq ! ( custom_data, expectation. custom_data) ;
199
213
200
214
let response = match msg {
201
215
TestCustomMessage :: Ping => TestCustomMessage :: Pong ,
0 commit comments