improve code structure

apps/web_app/telegram/crud.py

@@ -110,6 +110,12 @@ async def get_objects_by_filter(
             return await db.scalars(stmp).all()
 
     async def write_to_db(self, obj: ModelType) -> None:
+        """
+        Write an object to the database.
+
+        Args:
+            obj (ModelType): The object to be added to the database.
+        """
         async with self.Session() as db:
             db.add(obj)
             await db.commit()

apps/web_app/telegram/handlers/create_notification.py

@@ -3,30 +3,31 @@
 from aiogram.fsm.state import State, StatesGroup
 from database.models import NotificationData
 from telegram.crud import TelegramCrud
-from utils.values import ProtocolIDs
-from fastapi import status
+from .utils import kb
 
 create_notification_router = Router()
 
 
 class NotificationFormStates(StatesGroup):
     """States for the notification form process."""
+
     wallet_id = State()
     health_ratio_level = State()
     protocol_id = State()
 
 
+# Define constants for health ratio limits
+HEALTH_RATIO_MIN = 0
+HEALTH_RATIO_MAX = 10
+
+
 @create_notification_router.callback_query(F.data == "create_subscription")
 async def start_form(callback: types.CallbackQuery, state: FSMContext):
     """Initiates the notification creation form by asking for the wallet ID."""
     await state.set_state(NotificationFormStates.wallet_id)
-    await callback.message.edit_text(
+    return callback.message.edit_text(
         "Please enter your wallet ID:",
-        reply_markup=types.InlineKeyboardMarkup(
-            inline_keyboard=[
-                [types.InlineKeyboardButton(text="Cancel", callback_data="cancel_form")]
-            ]
-        ),
+        reply_markup=kb.cancel_form(),
     )
 
 
@@ -35,13 +36,9 @@ async def process_wallet_id(message: types.Message, state: FSMContext):
     """Processes the wallet ID input from the user."""
     await state.update_data(wallet_id=message.text)
     await state.set_state(NotificationFormStates.health_ratio_level)
-    await message.answer(
+    return message.answer(
         "Please enter your health ratio level (between 0 and 10):",
-        reply_markup=types.InlineKeyboardMarkup(
-            inline_keyboard=[
-                [types.InlineKeyboardButton(text="Cancel", callback_data="cancel_form")]
-            ]
-        ),
+        reply_markup=kb.cancel_form(),
     )
 
 
@@ -50,44 +47,22 @@ async def process_health_ratio(message: types.Message, state: FSMContext):
     """Processes the health ratio level input from the user."""
     try:
         health_ratio = float(message.text)
-        if not (0 <= health_ratio <= 10):
+        if not (HEALTH_RATIO_MIN <= health_ratio <= HEALTH_RATIO_MAX):
             raise ValueError
-
-        await state.update_data(health_ratio_level=health_ratio)
-        await state.set_state(NotificationFormStates.protocol_id)
-
-        # Create protocol selection buttons
-        protocol_buttons = []
-        for protocol in ProtocolIDs:
-            protocol_buttons.append(
-                [
-                    types.InlineKeyboardButton(
-                        text=protocol.value, callback_data=f"protocol_{protocol.value}"
-                    )
-                ]
-            )
-        protocol_buttons.append(
-            [types.InlineKeyboardButton(text="Cancel", callback_data="cancel_form")]
-        )
-
-        await message.answer(
-            "Please select your protocol:",
-            reply_markup=types.InlineKeyboardMarkup(inline_keyboard=protocol_buttons),
-        )
     except ValueError:
-        await message.answer(
+        return message.answer(
             "Please enter a valid number between 0 and 10.",
-            reply_markup=types.InlineKeyboardMarkup(
-                inline_keyboard=[
-                    [
-                        types.InlineKeyboardButton(
-                            text="Cancel", callback_data="cancel_form"
-                        )
-                    ]
-                ]
-            ),
+            reply_markup=kb.cancel_form(),
         )
 
+    await state.update_data(health_ratio_level=health_ratio)
+    await state.set_state(NotificationFormStates.protocol_id)
+
+    return message.answer(
+        "Please select your protocol:",
+        reply_markup=kb.protocols(),
+    )
+
 
 @create_notification_router.callback_query(F.data.startswith("protocol_"))
 async def process_protocol(
@@ -108,25 +83,17 @@ async def process_protocol(
     await crud.write_to_db(subscription)
 
     await state.clear()
-    await callback.message.edit_text(
+    return callback.message.edit_text(
         "Subscription created successfully!",
-        reply_markup=types.InlineKeyboardMarkup(
-            inline_keyboard=[
-                [types.InlineKeyboardButton(text="Go to menu", callback_data="go_menu")]
-            ]
-        ),
+        reply_markup=kb.go_menu(),
     )
 
 
 @create_notification_router.callback_query(F.data == "cancel_form")
 async def cancel_form(callback: types.CallbackQuery, state: FSMContext):
     """Cancels the form and clears the state."""
     await state.clear()
-    await callback.message.edit_text(
+    return callback.message.edit_text(
         "Form cancelled.",
-        reply_markup=types.InlineKeyboardMarkup(
-            inline_keyboard=[
-                [types.InlineKeyboardButton(text="Go to menu", callback_data="go_menu")]
-            ]
-        ),
+        reply_markup=kb.go_menu(),
     )

apps/web_app/telegram/handlers/utils/kb.py

@@ -2,6 +2,7 @@
 
 from aiogram.types import InlineKeyboardButton, InlineKeyboardMarkup
 from aiogram.utils.keyboard import InlineKeyboardBuilder
+from utils.values import ProtocolIDs
 
 
 def go_menu():
@@ -103,3 +104,21 @@ def pagination_notifications(curent_uuid: UUID, page: int):
     if page == 0:
         markup.adjust(2, 1, 1)
     return markup.as_markup()
+
+def cancel_form():
+    """
+    Returns an InlineKeyboardMarkup with a single button labeled "Cancel" with the callback data "cancel_form".
+    """
+    return InlineKeyboardMarkup(inline_keyboard=[[InlineKeyboardButton(text="Cancel", callback_data="cancel_form")]])
+
+
+def protocols():
+    """
+    Returns an InlineKeyboardMarkup with buttons for each protocol.
+    """
+    # Create protocol selection buttons
+    markup = InlineKeyboardBuilder()
+    for protocol in ProtocolIDs:
+        markup.button(text=protocol.name, callback_data=f"protocol_{protocol.value}")
+    markup.button(text="Cancel", callback_data="cancel_form")
+    return markup.as_markup()