def test_edits_update(self):
storage = self.get_storage()
bot = Bot()
context = Context(bot)
pin_handler = handlers.pinned(storage)
edit_handler = handlers.message_edited(storage)
message_amount = 5
msgs = gen_same_chat_messages(message_amount)
upds = [Update(msg, None) for msg in msgs]
msg1, msg2 = choice(msgs), choice(msgs)
for update, number in zip(upds, range(message_amount)):
pin_handler(update, context)
already_edited = message_amount - 1
self.assertEqual(len(bot.edited), already_edited)
upd1 = Update(msg1, None, msg1)
upd2 = Update(msg2, None, msg2)
edit_handler(upd1, context)
self.assertEqual(len(bot.edited), already_edited + 1)
edit_handler(upd2, context)
self.assertEqual(len(bot.edited), already_edited + 2)
edit_handler(upd1, context)
self.assertEqual(len(bot.edited), already_edited + 3)
def test_deletes_on_nothing(self):
storage = self.get_storage()
bot = Bot()
context = Context(bot)
pin_handler = handlers.pinned(storage)
button_handler = handlers.button_pressed(storage)
message_amount = 5
msgs = gen_same_chat_messages(message_amount)
msg_upds = [Update(msg, None) for msg in msgs]
unpins = [gen_unpin_data(msg) for msg in msgs]
unpin_upds = [Update(None, unpin) for unpin in unpins]
chat_id = msgs[0].chat.id
for pin_update in msg_upds:
pin_handler(pin_update, context)
for unpin_update in unpin_upds:
button_handler(unpin_update, context)
self.assertEqual(len(bot.deleted), 1)
self.assertNotEqual(len(bot.sent), 0)
self.assertEqual(bot.sent[0]['m_id'], bot.deleted[0]['m_id'])
self.assertEqual(bot.sent[0]['chat_id'], bot.deleted[0]['chat_id'])
sent_first_batch = len(bot.sent)
button_handler(unpin_upds[0], context)
self.assertEqual(len(bot.deleted), 1)
self.assertEqual(len(bot.sent), sent_first_batch)
def test_gathers_correct_links(self):
storage = self.get_storage()
bot = Bot()
context = Context(bot)
pin_handler = handlers.pinned(storage)
msg = gen_message()
link1 = "github.com"
link2 = "https://kde.org/"
start1 = 0
length1 = len(link1)
start2 = len(link1) + 1
length2 = len(link2)
msg.entities = [Entity(start1, length1), Entity(start2, length2)]
msg.text = "\n".join([link1, link2])
update = Update(msg, None)
pin_handler(update, context)
sent = bot.sent[-1]["text"]
link_re = re.compile('<a href="([^"]+)">')
links = link_re.findall(sent)
self.assertEqual(links[0], link1)
self.assertEqual(links[1], link2)
def test_handlers_store(self):
storage = self.get_storage()
bot = Bot()
context = Context(bot)
pin_handler = handlers.pinned(storage)
button_handler = handlers.button_pressed(storage)
message_amount = 5
msgs = gen_same_chat_messages(message_amount)
msg_upds = [Update(msg, None) for msg in msgs]
unpins = [gen_unpin_data(msg) for msg in msgs]
unpin_upds = [Update(None, unpin) for unpin in unpins]
chat_id = msgs[0].chat.id
it = zip(msg_upds, unpin_upds, range(1, message_amount + 1))
for pin_update, unpin_update, amount in it:
pin_handler(pin_update, context)
self.assertEqual(len(storage.get(chat_id)), amount)
button_handler(unpin_update, context)
self.assertEqual(len(storage.get(chat_id)), amount - 1)
# test deleting non-existant
button_handler(unpin_update, context)
self.assertEqual(len(storage.get(chat_id)), amount - 1)
#second add of deleted to keep amount increasing
pin_handler(pin_update, context)
self.assertEqual(len(storage.get(chat_id)), amount)
unpin_all_cb = Update.CbQuery(msgs[0], handlers.UnpinAll)
unpin_all_upd = Update(None, unpin_all_cb)
button_handler(unpin_all_upd, context)
self.assertFalse(storage.has(chat_id))
def run_example(msg_update, msg_context) -> None:
print("Got example command")
storage = LocalStorage()
bot = test.Bot()
context = test.Context(bot)
messages = []
# message with links
for _ in range(randint(5, 7)):
msg = test.gen_message(60 * 60 * 24 * 30)
if choice([True, False]):
# link message
link = choice(["github.com", "https://kde.org/"])
text = f"foo {link} bar"
start = 4
length = len(link)
msg.entities = [test.Entity(start, length)]
msg.text = text
else:
# regular message
pass
messages.append(msg)
# set correct chat_id
for msg in messages:
msg.chat = msg_update.message.chat
# do pinning
for msg in messages:
update = test.Update(msg, None)
handlers.pinned(storage)(update, context)
msg_context.bot.send_message(msg_update.message.chat_id,
text=bot.edited[-1]["text"],
parse_mode="HTML",
reply_markup=bot.edited[-1]["markup"])
print(bot.edited[-1])
def test_pin_sends_and_edits(self):
storage = self.get_storage()
bot = Bot()
context = Context(bot)
pin_handler = handlers.pinned(storage)
message_amount = 5
msgs = gen_same_chat_messages(message_amount)
upds = [Update(msg, None) for msg in msgs]
for update in upds:
pin_handler(update, context)
self.assertEqual(len(bot.pinned), message_amount)
self.assertEqual(len(bot.sent), 1)
self.assertEqual(len(bot.edited), message_amount - 1)
def test_keep_last(self):
storage = self.get_storage()
bot = Bot()
context = Context(bot)
pin_handler = handlers.pinned(storage)
button_handler = handlers.button_pressed(storage)
message_amount = 5
msgs = gen_same_chat_messages(message_amount)
upds = [Update(msg, None) for msg in msgs]
chat_id = msgs[0].chat.id
for update, amount in zip(upds, range(1, message_amount + 1)):
pin_handler(update, context)
self.assertEqual(len(storage.get(chat_id)), amount)
keep_last_cb = Update.CbQuery(msgs[0], handlers.KeepLast)
keep_last_upd = Update(None, keep_last_cb)
button_handler(keep_last_upd, context)
self.assertTrue(storage.has(chat_id))
self.assertEqual(len(storage.get(chat_id)), 1)
def test_user_message_resends(self):
storage = self.get_storage()
bot = Bot()
context = Context(bot)
pin_handler = handlers.pinned(storage)
message_handler = handlers.message(storage)
message_amount = 5
msgs = gen_same_chat_messages(message_amount)
upds = [Update(msg, None) for msg in msgs]
user_message = gen_message()
user_message.chat.id = msgs[0].chat.id
user_update = Update(user_message, None)
for update in upds:
pin_handler(update, context)
message_handler(user_update, context)
self.assertEqual(len(bot.pinned), message_amount)
self.assertEqual(len(bot.sent), message_amount)
self.assertEqual(len(bot.edited), 0)
def main(token: str) -> None:
updater = Updater(token, use_context=True)
dp = updater.dispatcher
storage: Union[Storage, LocalStorage] = Storage()
if "local" in sys.argv:
storage = LocalStorage()
print("Running with local storage")
# mundane handlers
dp.add_handler(CommandHandler("start", handlers.start))
dp.add_handler(CommandHandler("help", handlers.help))
# catch messages pinned
pin_filter = Filters.status_update.pinned_message
dp.add_handler(MessageHandler(pin_filter, handlers.pinned(storage)))
# catch presses of "unpin" buttons
dp.add_handler(CallbackQueryHandler(handlers.button_pressed(storage)))
# catch edited messages
edit_filter = Filters.update.edited_message
edit_handler = MessageHandler(edit_filter,
handlers.message_edited(storage))
dp.add_handler(edit_handler)
# catch any user message
msg_filter = ~Filters.status_update
dp.add_handler(MessageHandler(msg_filter, handlers.message(storage)))
# Enable logging
logging.basicConfig(
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
level=logging.INFO)
logger = logging.getLogger(__name__)
dp.add_error_handler(handlers.error(logger))
updater.start_polling()
updater.idle()
def test_resending_deletes_old(self):
storage = self.get_storage()
bot = Bot()
context = Context(bot)
pin_handler = handlers.pinned(storage)
message_handler = handlers.message(storage)
message_amount = 5
msgs = gen_same_chat_messages(message_amount)
upds = [Update(msg, None) for msg in msgs]
user_message = gen_message()
user_message.chat.id = msgs[0].chat.id
user_update = Update(user_message, None)
for update in upds:
pin_handler(update, context)
sent = bot.sent[-1]
message_handler(user_update, context)
pin_handler(update, context)
deleted = bot.deleted[-1]
self.assertEqual(sent['m_id'], deleted['m_id'])
self.assertEqual(sent['chat_id'], deleted['chat_id'])