Add support for UpdateKind::EditedMessage

teloxide_tests/src/dataset/message.rs

@@ -60,7 +60,7 @@ macro_rules! Message {
         }
 
         impl crate::dataset::IntoUpdate for $name {
-            /// Converts the MockCallbackQuery into an updates vector
+            /// Converts the mock message into an updates vector
             ///
             /// # Example
             /// ```
@@ -87,6 +87,59 @@ macro_rules! Message {
 
 pub(crate) use Message;
 
+#[derive(Clone, Debug, PartialEq)]
+pub struct MockEditedMessage(Message);
+
+impl MockEditedMessage {
+    /// Creates a new MockEditedMessage wrapper.
+    ///
+    /// This is useful for testing the `UpdateKind::EditedMessage` variant.
+    ///
+    /// # Example
+    /// ```
+    /// use chrono::Utc;
+    ///
+    /// let message = teloxide_tests::MockMessageText::new().edit_date(Utc::now()).build();
+    /// let edited_message = teloxide_tests::MockEditedMessage::new(message.clone());
+    /// assert_eq!(edited_message.message(), &message);
+    pub fn new(mut message: Message) -> Self {
+        if let MessageKind::Common(ref mut common) = message.kind {
+            common.edit_date = common.edit_date.or(Some(Utc::now()));
+        }
+        Self(message)
+    }
+
+    pub fn message(&self) -> &Message {
+        &self.0
+    }
+}
+
+impl crate::dataset::IntoUpdate for MockEditedMessage {
+    /// Converts the edited Message into an updates vector
+    ///
+    /// # Example
+    /// ```
+    /// use chrono::Utc;
+    /// use teloxide_tests::IntoUpdate;
+    /// use teloxide::types::{UpdateId, UpdateKind};
+    /// use std::sync::atomic::AtomicI32;
+    ///
+    /// let message = teloxide_tests::MockMessageText::new().edit_date(Utc::now()).build();
+    /// let edited_message = teloxide_tests::MockEditedMessage::new(message.clone());
+    /// let update = edited_message.into_update(&AtomicI32::new(42))[0].clone();
+    ///
+    /// assert_eq!(update.id, UpdateId(42));
+    /// assert_eq!(update.kind, UpdateKind::EditedMessage(message));
+    /// ```
+    ///
+    fn into_update(self, id: &AtomicI32) -> Vec<Update> {
+        vec![Update {
+            id: UpdateId(id.fetch_add(1, Ordering::Relaxed) as u32),
+            kind: UpdateKind::EditedMessage(self.0),
+        }]
+    }
+}
+
 // More messages like Webapp data is needed
 
 Message! {

teloxide_tests/src/mock_bot.rs

@@ -237,6 +237,10 @@ where
                     self.state.lock().unwrap().add_message(&mut message);
                     update.kind = UpdateKind::Message(message.clone());
                 }
+                UpdateKind::EditedMessage(mut message) => {
+                    self.state.lock().unwrap().edit_message(&mut message);
+                    update.kind = UpdateKind::EditedMessage(message.clone());
+                }
                 UpdateKind::CallbackQuery(mut callback) => {
                     if let Some(MaybeInaccessibleMessage::Regular(ref mut message)) =
                         callback.message

teloxide_tests/src/server/messages.rs

@@ -14,7 +14,20 @@ impl Messages {
         self.last_message_id
     }
 
-    pub fn edit_message<T>(&mut self, message_id: i32, field: &str, value: T) -> Option<Message>
+    pub fn edit_message(&mut self, message: Message) -> Option<Message> {
+        self.messages.iter().find(|m| m.id == message.id)?; // Find the message (return None if not found)
+
+        self.messages.retain(|m| m.id != message.id); // Remove the old message
+        self.messages.push(message.clone()); // Add the new message
+        Some(message) // Profit!
+    }
+
+    pub fn edit_message_field<T>(
+        &mut self,
+        message_id: i32,
+        field: &str,
+        value: T,
+    ) -> Option<Message>
     where
         T: Serialize,
     {
@@ -38,11 +51,11 @@ impl Messages {
         match reply_markup {
             None => {
                 // Telegram deletes reply markup when `editMessageText` is called without any.
-                self.edit_message(message_id, "reply_markup", None::<()>)
+                self.edit_message_field(message_id, "reply_markup", None::<()>)
             }
             // Only the inline keyboard can be inside of a message
             Some(ReplyMarkup::InlineKeyboard(reply_markup)) => {
-                self.edit_message(message_id, "reply_markup", reply_markup)
+                self.edit_message_field(message_id, "reply_markup", reply_markup)
             }
             _ => unreachable!("Only InlineKeyboard is allowed"),
         }
@@ -83,6 +96,7 @@ impl Messages {
 
 #[cfg(test)]
 mod tests {
+    use chrono::{TimeZone, Utc};
     use serial_test::serial;
     use teloxide::types::{InlineKeyboardButton, InlineKeyboardMarkup, MessageId};
 
@@ -116,15 +130,38 @@ mod tests {
 
     #[test]
     #[serial]
-    fn test_edit_messages() {
+    fn test_edit_message() {
+        let mut messages = Messages::default();
+        let date = Utc.with_ymd_and_hms(2020, 1, 1, 0, 0, 0).unwrap();
+        messages.add_message(
+            message_common::MockMessageText::new()
+                .text("123")
+                .id(1)
+                .build(),
+        );
+        messages.edit_message(
+            message_common::MockMessageText::new()
+                .text("321")
+                .edit_date(date)
+                .id(1)
+                .build(),
+        );
+        let message = messages.get_message(1).unwrap();
+        assert_eq!(message.text().unwrap(), "321");
+        assert_eq!(message.edit_date().unwrap(), &date);
+    }
+
+    #[test]
+    #[serial]
+    fn test_edit_message_field() {
         let mut messages = Messages::default();
         messages.add_message(
             message_common::MockMessageText::new()
                 .text("123")
                 .id(1)
                 .build(),
         );
-        messages.edit_message(1, "text", "1234");
+        messages.edit_message_field(1, "text", "1234");
         assert_eq!(messages.get_message(1).unwrap().text().unwrap(), "1234");
     }
 

teloxide_tests/src/server/routes/edit_message_caption.rs

@@ -35,8 +35,8 @@ pub async fn edit_message_caption(
             let mut lock = state.lock().unwrap();
             check_if_message_exists!(lock, message_id);
             lock.messages
-                .edit_message(message_id, "caption", body.caption.clone());
-            lock.messages.edit_message(
+                .edit_message_field(message_id, "caption", body.caption.clone());
+            lock.messages.edit_message_field(
                 message_id,
                 "caption_entities",
                 body.caption_entities.clone().unwrap_or_default(),

teloxide_tests/src/server/routes/edit_message_reply_markup.rs

@@ -37,11 +37,11 @@ pub async fn edit_message_reply_markup(
             let message = match body.reply_markup.clone() {
                 Some(reply_markup) => lock
                     .messages
-                    .edit_message(message_id, "reply_markup", reply_markup)
+                    .edit_message_field(message_id, "reply_markup", reply_markup)
                     .unwrap(),
                 None => lock
                     .messages
-                    .edit_message(message_id, "reply_markup", None::<()>)
+                    .edit_message_field(message_id, "reply_markup", None::<()>)
                     .unwrap(),
             };
 

teloxide_tests/src/server/routes/edit_message_text.rs

@@ -36,8 +36,8 @@ pub async fn edit_message_text(
             check_if_message_exists!(lock, message_id);
 
             lock.messages
-                .edit_message(message_id, "text", body.text.clone());
-            lock.messages.edit_message(
+                .edit_message_field(message_id, "text", body.text.clone());
+            lock.messages.edit_message_field(
                 message_id,
                 "entities",
                 body.entities.clone().unwrap_or(vec![]),

teloxide_tests/src/state.rs

@@ -50,4 +50,32 @@ impl State {
         log::debug!("Inserted message with {}.", message.id);
         self.messages.add_message(message.clone());
     }
+
+    pub(crate) fn edit_message(&mut self, message: &mut Message) {
+        let old_message = self.messages.get_message(message.id.0);
+
+        if old_message.is_none() {
+            log::error!(
+                "Not editing message with id {}, this id does not exist in the database.",
+                message.id
+            );
+            return;
+        }
+
+        if let Some(file_meta) = find_file(serde_json::to_value(&message).unwrap()) {
+            if self
+                .files
+                .iter()
+                .all(|f| f.meta.unique_id != file_meta.unique_id)
+            {
+                let file = File {
+                    meta: file_meta,
+                    path: "some_path.txt".to_string(), // This doesn't really matter
+                };
+                self.files.push(file);
+            }
+        }
+        log::debug!("Edited message with {}.", message.id);
+        self.messages.edit_message(message.clone());
+    }
 }

teloxide_tests/src/tests.rs

@@ -4,6 +4,7 @@ use std::{
     thread,
 };
 
+use chrono::Utc;
 use futures_util::future::BoxFuture;
 use serde::{Deserialize, Serialize};
 use teloxide::{
@@ -519,7 +520,11 @@ fn get_schema() -> UpdateHandler<Box<dyn std::error::Error + Send + Sync + 'stat
                 .filter_command::<AllCommands>()
                 .endpoint(handler),
         )
-        .branch(Update::filter_message().endpoint(handler))
+        .branch(
+            Update::filter_edited_message()
+                .filter_command::<AllCommands>()
+                .branch(case![AllCommands::ForwardMessage].endpoint(handler)),
+        )
         .branch(Update::filter_callback_query().endpoint(callback_handler))
 }
 
@@ -1088,6 +1093,8 @@ async fn test_forward_message() {
         last_sent_message.forward_from_user().unwrap().id,
         first_sent_message.from.as_ref().unwrap().id
     );
+    assert_eq!(responses.forwarded_messages.len(), 1);
+    assert_eq!(&responses.forwarded_messages[0].message, last_sent_message);
 }
 
 #[tokio::test]
@@ -1200,3 +1207,29 @@ async fn test_send_invoice() {
         "Absolutely Nothing"
     );
 }
+
+#[tokio::test]
+async fn test_edited_message() {
+    let mock_message = MockMessageText::new().text("/forwardmessage first");
+    let mut bot = MockBot::new(mock_message.clone(), get_schema());
+    bot.dispatch().await;
+
+    let responses = bot.get_responses();
+    assert_eq!(responses.forwarded_messages.len(), 1);
+    let forwarded_message = &responses.forwarded_messages[0].message;
+    assert_eq!(forwarded_message.text(), Some("/forwardmessage first"));
+
+    let edited_message = MockEditedMessage::new(
+        mock_message
+            .text("/forwardmessage second")
+            .edit_date(Utc::now())
+            .build(),
+    );
+    bot.update(edited_message);
+    bot.dispatch().await;
+
+    let responses = bot.get_responses();
+    assert_eq!(responses.forwarded_messages.len(), 1);
+    let forwarded_message = &responses.forwarded_messages[0].message;
+    assert_eq!(forwarded_message.text(), Some("/forwardmessage second"));
+}