Fix manage_worker

Author: reo7sp

tests/test_manage_worker.py

@@ -39,10 +39,10 @@ def test_does_nothing_if_state_not_exists(self):
 
 class TestStart:
     def test_start_sends_welcome_message(self):
-        bot = Mock()
         update = Mock()
+        context = Mock()
 
-        start(bot, update)
+        start(update, context)
 
         update.message.reply_text.assert_called_once()
         call_args = update.message.reply_text.call_args[0][0]
@@ -52,24 +52,24 @@ def test_start_sends_welcome_message(self):
 class TestNew:
     @patch('vk_channelify.manage_worker.metrics')
     def test_new_starts_conversation(self, mock_metrics):
-        bot = Mock()
         update = Mock()
+        context = Mock()
 
-        result = new(bot, update)
+        result = new(update, context)
 
         assert_that(result, equal_to(ASKED_VK_GROUP_LINK_IN_NEW))
         update.message.reply_text.assert_called_once()
 
 
 class TestNewInStateAskedVkGroupLink:
     def test_saves_vk_domain_and_asks_for_channel_access(self):
-        bot = Mock()
         update = Mock()
+        context = Mock()
         update.message.text = 'https://vk.ru/mygroup'
         update.message.from_user.id = 12345
         users_state = {}
 
-        result = new_in_state_asked_vk_group_link(bot, update, users_state=users_state)
+        result = new_in_state_asked_vk_group_link(update, context, users_state=users_state)
 
         assert_that(result, equal_to(ASKED_CHANNEL_ACCESS_IN_NEW))
         assert_that(users_state[12345]['vk_domain'], equal_to('mygroup'))
@@ -78,10 +78,10 @@ def test_saves_vk_domain_and_asks_for_channel_access(self):
 
 class TestNewInStateAskedChannelAccess:
     def test_asks_for_channel_message(self):
-        bot = Mock()
         update = Mock()
+        context = Mock()
 
-        result = new_in_state_asked_channel_access(bot, update)
+        result = new_in_state_asked_channel_access(update, context)
 
         assert_that(result, equal_to(ASKED_CHANNEL_MESSAGE_IN_NEW))
         update.message.reply_text.assert_called_once()
@@ -92,27 +92,27 @@ class TestNewInStateAskedChannelMessage:
     @patch('vk_channelify.manage_worker.Channel')
     @patch('vk_channelify.manage_worker.DisabledChannel')
     def test_creates_channel_successfully(self, mock_disabled_channel, mock_channel, mock_metrics):
-        bot = Mock()
         update = Mock()
+        context = Mock()
         update.message.from_user.id = 12345
         update.message.from_user.username = 'testuser'
         update.message.forward_from_chat.id = -100123456
         users_state = {12345: {'vk_domain': 'mygroup'}}
         db = Mock()
         db_session_maker = Mock(return_value=db)
 
-        result = new_in_state_asked_channel_message(bot, update, db_session_maker=db_session_maker, users_state=users_state)
+        result = new_in_state_asked_channel_message(update, context, db_session_maker=db_session_maker, users_state=users_state)
 
         assert_that(result, equal_to(ConversationHandler.END))
         db.add.assert_called_once()
         db.commit.assert_called_once()
-        bot.send_message.assert_called_once()
+        context.bot.send_message.assert_called_once()
 
     @patch('vk_channelify.manage_worker.metrics')
     @patch('vk_channelify.manage_worker.Channel')
     def test_rolls_back_on_error(self, mock_channel, mock_metrics):
-        bot = Mock()
         update = Mock()
+        context = Mock()
         update.message.from_user.id = 12345
         update.message.from_user.username = 'testuser'
         update.message.forward_from_chat.id = -100123456
@@ -122,7 +122,7 @@ def test_rolls_back_on_error(self, mock_channel, mock_metrics):
         db_session_maker = Mock(return_value=db)
 
         with pytest.raises(Exception):
-            new_in_state_asked_channel_message(bot, update, db_session_maker=db_session_maker, users_state=users_state)
+            new_in_state_asked_channel_message(update, context, db_session_maker=db_session_maker, users_state=users_state)
 
         db.rollback.assert_called_once()
         mock_metrics.telegram_conversations_total.labels.assert_called_with(type='new', status='failed')
@@ -131,12 +131,12 @@ def test_rolls_back_on_error(self, mock_channel, mock_metrics):
 class TestCancelNew:
     @patch('vk_channelify.manage_worker.metrics')
     def test_cancel_ends_conversation(self, mock_metrics):
-        bot = Mock()
         update = Mock()
+        context = Mock()
         update.message.from_user.id = 12345
         users_state = {12345: {'vk_domain': 'mygroup'}}
 
-        result = cancel_new(bot, update, users_state=users_state)
+        result = cancel_new(update, context, users_state=users_state)
 
         assert_that(result, equal_to(ConversationHandler.END))
         assert_that(12345 not in users_state, is_(True))

vk_channelify/manage_worker.py

@@ -102,9 +102,9 @@ def on_error(update, context):
 
 def catch_exceptions(func):
     @wraps(func)
-    def wrapper(bot, update, *args, **kwargs):
+    def wrapper(update, context, *args, **kwargs):
         try:
-            return func(bot, update, *args, **kwargs)
+            return func(update, context, *args, **kwargs)
         except Exception as e:
             logger.error('Exception in {}: {}'.format(func.__name__, e))
             traceback.print_exc()
@@ -116,11 +116,11 @@ def wrapper(bot, update, *args, **kwargs):
 def observe_metrics(command_name):
     def decorator(func):
         @wraps(func)
-        def wrapper(bot, update, *args, **kwargs):
+        def wrapper(update, context, *args, **kwargs):
             metrics.telegram_commands_total.labels(command=command_name).inc()
             start_time = time.time()
             try:
-                result = func(bot, update, *args, **kwargs)
+                result = func(update, context, *args, **kwargs)
                 duration = time.time() - start_time
                 metrics.telegram_command_duration_seconds.labels(command=command_name).observe(duration)
                 return result
@@ -147,13 +147,13 @@ def wrapper(*args, db_session_maker, **kwargs):
 
 @catch_exceptions
 @observe_metrics('start')
-def start(bot, update):
+def start(update, context):
     update.message.reply_text('Команда /new настроит новый канал. В канал будут пересылаться посты из группы ВК')
 
 
 @catch_exceptions
 @observe_metrics('new')
-def new(bot, update):
+def new(update, context):
     metrics.telegram_conversations_total.labels(type='new', status='started').inc()
 
     update.message.reply_text('Отправьте ссылку на группу ВК')
@@ -162,7 +162,7 @@ def new(bot, update):
 
 
 @catch_exceptions
-def new_in_state_asked_vk_group_link(bot, update, users_state):
+def new_in_state_asked_vk_group_link(update, context, users_state):
     vk_url = update.message.text
     vk_domain = vk_url.split('/')[-1]
     users_state[update.message.from_user.id] = dict()
@@ -178,15 +178,15 @@ def new_in_state_asked_vk_group_link(bot, update, users_state):
 
 
 @catch_exceptions
-def new_in_state_asked_channel_access(bot, update):
+def new_in_state_asked_channel_access(update, context):
     update.message.reply_text('Хорошо. Перешлите любое сообщение из канала', reply_markup=ReplyKeyboardRemove())
 
     return ASKED_CHANNEL_MESSAGE_IN_NEW
 
 
 @catch_exceptions
 @make_db_session
-def new_in_state_asked_channel_message(bot, update, db, users_state):
+def new_in_state_asked_channel_message(update, context, db, users_state):
     user_id = update.message.from_user.id
     username = update.message.from_user.username
     channel_id = str(update.message.forward_from_chat.id)
@@ -208,7 +208,7 @@ def new_in_state_asked_channel_message(bot, update, db, users_state):
         logger.warning('Cannot delete disabled channel of {}'.format(channel_id))
         traceback.print_exc()
 
-    bot.send_message(channel_id, 'Канал работает с помощью @vk_channelify_bot')
+    context.bot.send_message(channel_id, 'Канал работает с помощью @vk_channelify_bot')
 
     update.message.reply_text('Готово!')
     update.message.reply_text('Бот будет проверять группу каждые 15 минут')
@@ -221,7 +221,7 @@ def new_in_state_asked_channel_message(bot, update, db, users_state):
 
 
 @catch_exceptions
-def cancel_new(bot, update, users_state):
+def cancel_new(update, context, users_state):
     metrics.telegram_conversations_total.labels(type='new', status='cancelled').inc()
 
     update.message.reply_text('Ладно', reply_markup=ReplyKeyboardRemove())
@@ -235,7 +235,7 @@ def cancel_new(bot, update, users_state):
 @catch_exceptions
 @make_db_session
 @observe_metrics('filter_by_hashtag')
-def filter_by_hashtag(bot, update, db, users_state):
+def filter_by_hashtag(update, context, db, users_state):
     user_id = update.message.from_user.id
 
     metrics.telegram_conversations_total.labels(type='filter_by_hashtag', status='started').inc()
@@ -246,7 +246,7 @@ def filter_by_hashtag(bot, update, db, users_state):
     keyboard_row = []
     for channel in db.query(Channel).filter(Channel.owner_id == str(user_id)).order_by(Channel.created_at.desc()):
         try:
-            channel_chat = bot.get_chat(chat_id=channel.channel_id)
+            channel_chat = context.bot.get_chat(chat_id=channel.channel_id)
             users_state[user_id]['channels'][channel_chat.title] = channel.channel_id
             keyboard_row.append(channel_chat.title)
             if len(keyboard_row) == 2:
@@ -265,7 +265,7 @@ def filter_by_hashtag(bot, update, db, users_state):
 
 @catch_exceptions
 @make_db_session
-def filter_by_hashtag_in_state_asked_channel_id(bot, update, db, users_state):
+def filter_by_hashtag_in_state_asked_channel_id(update, context, db, users_state):
     user_id = update.message.from_user.id
     channel_title = update.message.text
     channel_id = str(users_state[user_id]['channels'][channel_title])
@@ -282,7 +282,7 @@ def filter_by_hashtag_in_state_asked_channel_id(bot, update, db, users_state):
 
 @catch_exceptions
 @make_db_session
-def filter_by_hashtag_in_state_asked_hashtags(bot, update, db, users_state):
+def filter_by_hashtag_in_state_asked_hashtags(update, context, db, users_state):
     user_id = update.message.from_user.id
     channel = users_state[user_id]['channel']
 
@@ -303,8 +303,9 @@ def filter_by_hashtag_in_state_asked_hashtags(bot, update, db, users_state):
 
 
 @catch_exceptions
-def cancel_filter_by_hashtag(bot, update, users_state):
+def cancel_filter_by_hashtag(update, context, users_state):
     metrics.telegram_conversations_total.labels(type='filter_by_hashtag', status='cancelled').inc()
+
     update.message.reply_text('Ладно', reply_markup=ReplyKeyboardRemove())
     update.message.reply_text('Настроить фильтр по хештегам можно командой /filter_by_hashtag')
     update.message.reply_text('Команда /new настроит новый канал')
@@ -317,7 +318,7 @@ def cancel_filter_by_hashtag(bot, update, users_state):
 @catch_exceptions
 @make_db_session
 @observe_metrics('recover')
-def recover(bot, update, db, users_state):
+def recover(update, context, db, users_state):
     user_id = update.message.from_user.id
 
     metrics.telegram_conversations_total.labels(type='recover', status='started').inc()
@@ -349,7 +350,7 @@ def recover(bot, update, db, users_state):
 
 @catch_exceptions
 @make_db_session
-def recover_in_state_asked_channel_id(bot, update, db, users_state):
+def recover_in_state_asked_channel_id(update, context, db, users_state):
     user_id = update.message.from_user.id
     channel_title = update.message.text
     channel_id = str(users_state[user_id]['channels'][channel_title])
@@ -381,7 +382,7 @@ def recover_in_state_asked_channel_id(bot, update, db, users_state):
 
 
 @catch_exceptions
-def cancel_recover(bot, update, users_state):
+def cancel_recover(update, context, users_state):
     metrics.telegram_conversations_total.labels(type='recover', status='cancelled').inc()
 
     update.message.reply_text('Ладно', reply_markup=ReplyKeyboardRemove())