1# ruff: noqa: PT009
2import unittest
3
4from telegram import ReplyKeyboardMarkup
5from telegram.ext import CommandHandler, ConversationHandler, MessageHandler, RegexHandler, Updater, filters
6
7from ptbtest import ChatGenerator, MessageGenerator, Mockbot, UserGenerator
8
9"""
10This is an example to show how the ptbtest suite can be used.
11This example follows the conversationbot2 example at:
12https://github.com/python-telegram-bot/python-telegram-bot/blob/master/examples/conversationbot2.py
13
14"""
15
16
17class TestConversationbot2(unittest.TestCase):
18 def setUp(self):
19 self.bot = Mockbot()
20 self.cg = ChatGenerator()
21 self.ug = UserGenerator()
22 self.mg = MessageGenerator(self.bot)
23 self.updater = Updater(bot=self.bot)
24
25 def test_conversation(self):
26 choosing, typing_reply, typing_choice = range(3)
27
28 reply_keyboard = [["Age", "Favourite colour"], ["Number of siblings", "Something else..."], ["Done"]]
29 markup = ReplyKeyboardMarkup(reply_keyboard, one_time_keyboard=True)
30
31 def facts_to_str(user_data):
32 facts = []
33 for key, value in user_data.items():
34 facts.append(f"{key} - {value}")
35
36 return "\n".join(facts).join(["\n", "\n"])
37
38 def start(update):
39 update.message.reply_text(
40 "Hi! My name is Doctor Botter. I will hold a more complex conversation with you. "
41 "Why don't you tell me something about yourself?",
42 reply_markup=markup,
43 )
44
45 return choosing
46
47 def regular_choice(update, user_data):
48 text = update.message.text
49 user_data["choice"] = text
50 update.message.reply_text(f"Your {text.lower()}? Yes, I would love to hear about that!")
51
52 return typing_reply
53
54 def custom_choice(update):
55 update.message.reply_text(
56 "Alright, please send me the category first, " 'for example "Most impressive skill"'
57 )
58
59 return typing_choice
60
61 def received_information(update, user_data):
62 text = update.message.text
63 category = user_data["choice"]
64 user_data[category] = text
65 del user_data["choice"]
66
67 update.message.reply_text(
68 f"Neat! Just so you know, this is what you already told me:"
69 f"{facts_to_str(user_data)}"
70 f"You can tell me more, or change your opinion on something.",
71 reply_markup=markup,
72 )
73
74 return choosing
75
76 def done(update, user_data):
77 if "choice" in user_data:
78 del user_data["choice"]
79
80 update.message.reply_text(
81 f"I learned these facts about you:" f"{facts_to_str(user_data)}" f"Until next time!"
82 )
83
84 user_data.clear()
85 return ConversationHandler.END
86
87 conv_handler = ConversationHandler(
88 entry_points=[CommandHandler("start", start)],
89 states={
90 choosing: [
91 RegexHandler("^(Age|Favourite colour|Number of siblings)$", regular_choice, pass_user_data=True),
92 RegexHandler("^Something else...$", custom_choice),
93 ],
94 typing_choice: [
95 MessageHandler(filters.text, regular_choice, pass_user_data=True),
96 ],
97 typing_reply: [
98 MessageHandler(filters.text, received_information, pass_user_data=True),
99 ],
100 },
101 fallbacks=[RegexHandler("^Done$", done, pass_user_data=True)],
102 )
103 dp = self.updater.dispatcher
104 dp.add_handler(conv_handler)
105 self.updater.start_polling()
106
107 # We are going to test a conversationhandler. Since this is tied in with user and chat we need to
108 # create both for consistancy
109 user = self.ug.get_user()
110 chat = self.cg.get_chat(type="group")
111 user2 = self.ug.get_user()
112 chat2 = self.cg.get_chat(user=user)
113
114 # let's start the conversation
115 u = self.mg.get_message(user=user, chat=chat, text="/start")
116 self.bot.insert_update(u)
117 data = self.bot.sent_messages[-1]
118 self.assertRegex(data["text"], r"Doctor Botter\. I will")
119 u = self.mg.get_message(user=user, chat=chat, text="Age")
120 self.bot.insert_update(u)
121 data = self.bot.sent_messages[-1]
122 self.assertRegex(data["text"], r"Your age\? Yes")
123
124 # now let's see what happens when another user in another chat starts conversating with the bot
125 u = self.mg.get_message(user=user2, chat=chat2, text="/start")
126 self.bot.insert_update(u)
127 data = self.bot.sent_messages[-1]
128 self.assertRegex(data["text"], r"Doctor Botter\. I will")
129 self.assertEqual(data["chat_id"], chat2.id)
130 self.assertNotEqual(data["chat_id"], chat.id)
131 # and cancels his conv.
132 u = self.mg.get_message(user=user2, chat=chat2, text="Done")
133 self.bot.insert_update(u)
134 data = self.bot.sent_messages[-1]
135 self.assertRegex(data["text"], r"Until next time!")
136
137 # cary on with first user
138 u = self.mg.get_message(user=user, chat=chat, text="23")
139 self.bot.insert_update(u)
140 data = self.bot.sent_messages[-1]
141 self.assertRegex(data["text"], r"Age - 23")
142 u = self.mg.get_message(user=user, chat=chat, text="Something else...")
143 self.bot.insert_update(u)
144 data = self.bot.sent_messages[-1]
145 self.assertRegex(data["text"], r"Most impressive skill")
146 u = self.mg.get_message(user=user, chat=chat, text="programming skill")
147 self.bot.insert_update(u)
148 data = self.bot.sent_messages[-1]
149 self.assertRegex(data["text"], r"Your programming skill\? Yes")
150 u = self.mg.get_message(user=user, chat=chat, text="High")
151 self.bot.insert_update(u)
152 data = self.bot.sent_messages[-1]
153 self.assertRegex(data["text"], r"programming skill - High")
154 u = self.mg.get_message(user=user, chat=chat, text="Done")
155 self.bot.insert_update(u)
156 data = self.bot.sent_messages[-1]
157 self.assertRegex(data["text"], r"programming skill - High")
158 self.assertRegex(data["text"], r"Age - 23")
159
160 self.updater.stop()
161
162
163if __name__ == "__main__":
164 unittest.main()