- fn handle_custom_message(&self, _msg: Self::CustomMessage) {
- self.num_messages_expected.fetch_sub(1, Ordering::SeqCst);
+ fn handle_custom_message(&self, msg: Self::CustomMessage) -> Option<Self::CustomMessage> {
+ match self.expected_messages.lock().unwrap().pop_front() {
+ Some(expected_msg) => assert_eq!(expected_msg, msg),
+ None => panic!("Unexpected message: {:?}", msg),
+ }
+
+ match msg {
+ TestCustomMessage::Request => Some(TestCustomMessage::Response),
+ TestCustomMessage::Response => None,
+ }