From b16630098a99cd3122a3618eaad3a081c2108ba5 Mon Sep 17 00:00:00 2001
From: LasterAlex <sasha.ishukshin@gmail.com>
Date: Fri, 7 Mar 2025 18:53:49 +0200
Subject: [PATCH] Added mocking of update poll

Fmt
---
 teloxide_tests/src/dataset/mod.rs    | 12 +++-
 teloxide_tests/src/dataset/tests.rs  | 22 +++++++-
 teloxide_tests/src/dataset/update.rs | 82 ++++++++++++++++++++++++++++
 3 files changed, 114 insertions(+), 2 deletions(-)
 create mode 100644 teloxide_tests/src/dataset/update.rs

diff --git a/teloxide_tests/src/dataset/mod.rs b/teloxide_tests/src/dataset/mod.rs
index 2ccdde5..ef8d7e4 100644
--- a/teloxide_tests/src/dataset/mod.rs
+++ b/teloxide_tests/src/dataset/mod.rs
@@ -5,7 +5,7 @@ use mime::Mime;
 use proc_macros::Changeable;
 use teloxide::types::{
     ChatPhoto, FileMeta, LinkPreviewOptions, LivePeriod, Location, Me, PhotoSize, Seconds, Update,
-    User, UserId, Video,
+    UpdateId, User, UserId, Video,
 };
 pub mod chat;
 pub mod chat_full_info;
@@ -13,12 +13,14 @@ pub mod chat_full_info;
 pub mod message;
 pub mod message_common;
 pub mod queries;
+pub mod update;
 pub use chat::*;
 pub use chat_full_info::*;
 pub use message::*;
 pub use message_common::*;
 pub use queries::*;
 use teloxide_tests_macros as proc_macros;
+pub use update::*;
 #[cfg(test)]
 mod tests;
 
@@ -42,6 +44,14 @@ where
     }
 }
 
+// Just to be able to use raw updates anywhere
+impl IntoUpdate for Update {
+    fn into_update(mut self, id: &AtomicI32) -> Vec<Update> {
+        self.id = UpdateId(id.fetch_add(1, Ordering::Relaxed) as u32);
+        vec![self]
+    }
+}
+
 //
 //  Structs below are just misc mocked structs
 //
diff --git a/teloxide_tests/src/dataset/tests.rs b/teloxide_tests/src/dataset/tests.rs
index 4136cdb..915c41b 100644
--- a/teloxide_tests/src/dataset/tests.rs
+++ b/teloxide_tests/src/dataset/tests.rs
@@ -1,7 +1,8 @@
 use teloxide::{
     dispatching::dialogue::GetChatId,
-    types::{ChatId, MessageEntity, MessageId, UpdateId, UserId},
+    types::{ChatId, MessageEntity, MessageId, UpdateId, UpdateKind, UserId},
 };
+use update::MockUpdatePoll;
 
 use crate::{dataset::*, proc_macros::Changeable};
 
@@ -399,3 +400,22 @@ fn test_callback_query() {
     assert_eq!(query_object.id, MockCallbackQuery::ID);
     assert_eq!(query_object.from.first_name, MockUser::FIRST_NAME);
 }
+
+//
+//
+//
+
+#[test]
+fn test_update_poll() {
+    let update = MockUpdatePoll::new().poll_id("123");
+
+    let update_object = update.into_update(&AtomicI32::new(1))[0].clone();
+
+    if let UpdateKind::Poll(poll) = update_object.kind {
+        assert_eq!(poll.question, MockMessagePoll::QUESTION);
+        assert_eq!(poll.poll_type, MockMessagePoll::POLL_TYPE);
+        assert_eq!(poll.id, "123".to_owned());
+    } else {
+        unreachable!()
+    }
+}
diff --git a/teloxide_tests/src/dataset/update.rs b/teloxide_tests/src/dataset/update.rs
new file mode 100644
index 0000000..ce610dd
--- /dev/null
+++ b/teloxide_tests/src/dataset/update.rs
@@ -0,0 +1,82 @@
+use std::sync::atomic::Ordering;
+
+use chrono::{DateTime, Utc};
+use teloxide::types::{
+    MessageEntity, Poll, PollOption, PollType, Seconds, Update, UpdateId, UpdateKind,
+};
+use teloxide_tests_macros::Changeable;
+
+use super::{IntoUpdate, MockMessagePoll};
+
+#[derive(Changeable, Clone)]
+pub struct MockUpdatePoll {
+    pub poll_id: String,
+    pub question: String,
+    pub question_entities: Option<Vec<MessageEntity>>,
+    pub options: Vec<PollOption>,
+    pub is_closed: bool,
+    pub total_voter_count: u32,
+    pub is_anonymous: bool,
+    pub poll_type: PollType,
+    pub allows_multiple_answers: bool,
+    pub correct_option_id: Option<u8>,
+    pub explanation: Option<String>,
+    pub explanation_entities: Option<Vec<MessageEntity>>,
+    pub open_period: Option<Seconds>,
+    pub close_date: Option<DateTime<Utc>>,
+}
+
+impl MockUpdatePoll {
+    /// Creates a new easily changable poll update builder
+    ///
+    /// # Example
+    /// ```
+    /// let update = teloxide_tests::MockUpdatePoll::new()
+    ///     .poll_id("123456");
+    ///
+    /// assert_eq!(update.poll_id, "123456");
+    /// ```
+    pub fn new() -> Self {
+        let poll = MockMessagePoll::new();
+        Self {
+            poll_id: poll.poll_id,
+            question: poll.question,
+            question_entities: poll.question_entities,
+            options: poll.options,
+            is_closed: poll.is_closed,
+            total_voter_count: poll.total_voter_count,
+            is_anonymous: poll.is_anonymous,
+            poll_type: poll.poll_type,
+            allows_multiple_answers: poll.allows_multiple_answers,
+            correct_option_id: poll.correct_option_id,
+            explanation: poll.explanation,
+            explanation_entities: poll.explanation_entities,
+            open_period: poll.open_period,
+            close_date: poll.close_date,
+        }
+    }
+}
+
+impl IntoUpdate for MockUpdatePoll {
+    fn into_update(self, id: &std::sync::atomic::AtomicI32) -> Vec<Update> {
+        vec![Update {
+            id: UpdateId(id.fetch_add(1, Ordering::Relaxed) as u32),
+            kind: UpdateKind::Poll(Poll {
+                id: self.poll_id,
+                question: self.question,
+                question_entities: self.question_entities,
+                options: self.options,
+                is_closed: self.is_closed,
+                total_voter_count: self.total_voter_count,
+                is_anonymous: self.is_anonymous,
+                poll_type: self.poll_type,
+                allows_multiple_answers: self.allows_multiple_answers,
+                correct_option_id: self.correct_option_id,
+                explanation: self.explanation,
+                explanation_entities: self.explanation_entities,
+                open_period: self.open_period,
+                close_date: self.close_date,
+            }),
+        }]
+    }
+}