1# ruff: noqa: PT009
2import time
3import unittest
4
5from telegram.ext import CommandHandler, Job, Updater
6
7from ptbtest import ChatGenerator, MessageGenerator, Mockbot
8
9"""
10This is an example to show how the ptbtest suite can be used.
11This example follows the timerbot example at:
12https://github.com/python-telegram-bot/python-telegram-bot/blob/master/examples/timerbot.py
13We will skip the start and help handlers and focus on the timer.
14
15"""
16
17
18class Testtimerbot(unittest.TestCase):
19 def setUp(self):
20 # For use within the tests we nee some stuff. Starting with a Mockbot
21 self.bot = Mockbot()
22 # Some generators for users and chats
23 self.cg = ChatGenerator()
24 # And a Messagegenerator and updater (for use with the bot.)
25 self.mg = MessageGenerator(self.bot)
26 self.updater = Updater(bot=self.bot)
27
28 def test_timer(self):
29 # first declare the callback methods
30 def alarm(bot, job):
31 """Function to send the alarm message"""
32 bot.sendMessage(job.context, text="Beep!")
33
34 def set_alarm(update, args, job_queue, chat_data):
35 """Adds a job to the queue"""
36 chat_id = update.message.chat_id
37 try:
38 # args[0] should contain the time for the timer in seconds
39 due = int(args[0])
40 if due < 0:
41 update.message.reply_text("Sorry we can not go back to future!")
42 return
43
44 # Add job to queue
45 job = Job(alarm, due, repeat=False, context=chat_id)
46 chat_data["job"] = job
47 job_queue.put(job)
48
49 update.message.reply_text("Timer successfully set!")
50
51 except (IndexError, ValueError):
52 update.message.reply_text("Usage: /set <seconds>")
53
54 def unset_alarm(update, chat_data):
55 """Removes the job if the user changed their mind"""
56
57 if "job" not in chat_data:
58 update.message.reply_text("You have no active timer")
59 return
60
61 job = chat_data["job"]
62 job.schedule_removal()
63 del chat_data["job"]
64
65 update.message.reply_text("Timer successfully unset!")
66
67 # Now add those handlers to the updater and start polling
68 dp = self.updater.dispatcher
69 dp.add_handler(CommandHandler("set", set_alarm, pass_args=True, pass_job_queue=True, pass_chat_data=True))
70 dp.add_handler(CommandHandler("unset", unset_alarm, pass_chat_data=True))
71 self.updater.start_polling()
72
73 # we want to check if the bot returns a message to the same chat after a period of time
74 # so let's create a chat to use
75 chat = self.cg.get_chat()
76
77 # let's generate some updates we can use
78 u1 = self.mg.get_message(chat=chat, text="/set")
79 u2 = self.mg.get_message(chat=chat, text="/set -20")
80 u3 = self.mg.get_message(chat=chat, text="/set 6")
81 u4 = self.mg.get_message(chat=chat, text="/unset")
82
83 # first check some errors
84 self.bot.insert_update(u1)
85 self.bot.insert_update(u2)
86 self.bot.insert_update(u4)
87 data = self.bot.sent_messages
88 self.assertEqual(len(data), 3)
89 self.assertEqual(data[0]["text"], "Usage: /set <seconds>")
90 self.assertEqual(data[1]["text"], "Sorry we can not go back to future!")
91 self.assertEqual(data[2]["text"], "You have no active timer")
92
93 # now check if setting and unsetting works (within timer limit)
94 self.bot.insert_update(u3)
95 data = self.bot.sent_messages[-1]
96 self.assertEqual(data["text"], "Timer successfully set!")
97 time.sleep(2)
98 self.bot.insert_update(u4)
99 data = self.bot.sent_messages[-1]
100 self.assertEqual(data["text"], "Timer successfully unset!")
101 # and to be certain we have to wait some more to see if it stops sending the message
102 # we reset the bot so we can be sure nothing more has been sent
103 self.bot.reset()
104 time.sleep(5)
105 data = self.bot.sent_messages
106 self.assertEqual(len(data), 0)
107
108 # lastly we will make sure an alarm message is sent after the timelimit has passed
109 self.bot.insert_update(u3)
110 time.sleep(6)
111 data = self.bot.sent_messages[-1]
112 self.assertEqual(data["text"], "Beep!")
113 self.assertEqual(data["chat_id"], chat.id)
114
115 # and stop the updater.
116 self.updater.stop()
117
118
119if __name__ == "__main__":
120 unittest.main()