refactor: redesign notification module

.dockerignore

@@ -38,4 +38,5 @@ config
 .pytest_cache
 test
 .env
-test.py
+test.py
+/backend/src/test

.github/workflows/build.yml

@@ -20,8 +20,7 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          if [ -f backend/requirements.txt ]; then pip install -r backend/requirements.txt; fi
-          pip install pytest
+          if [ -f backend/requirements-dev.txt ]; then pip install -r backend/requirements-dev.txt; fi
       - name: Test
         working-directory: ./backend/src
         run: |

.gitignore

@@ -212,8 +212,3 @@ dev-dist
 
 # vitepress
 /docs/.vitepress/cache/
-
-
-# test file
-test.*
-test_*

backend/requirements-dev.txt

@@ -2,4 +2,10 @@
 ruff
 black
 pre-commit
-pytest
+
+# testing
+pytest          # testing framework
+pytest-httpx    # httpx mock testing
+pytest-asyncio  # asyncio testing
+pytest-cov      # testing coverage
+pytest-mock     # mock wrapper for pytest

backend/requirements.txt

@@ -27,3 +27,6 @@ sqlmodel==0.0.8
 sse-starlette==1.6.5
 semver==3.0.1
 openai==0.28.1
+tenacity==8.2.3
+litequeue==0.7
+httpx==0.25.0

backend/src/module/api/__init__.py

@@ -4,6 +4,7 @@
 from .bangumi import router as bangumi_router
 from .config import router as config_router
 from .log import router as log_router
+from .notification import router as notification_router
 from .program import router as program_router
 from .rss import router as rss_router
 from .search import router as search_router
@@ -19,3 +20,4 @@
 v1.include_router(config_router)
 v1.include_router(rss_router)
 v1.include_router(search_router)
+v1.include_router(notification_router)

backend/src/module/api/notification.py

@@ -0,0 +1,132 @@
+from fastapi import APIRouter, Depends, HTTPException, Query
+
+from module.conf.config import settings
+from module.models.api import (
+    NotificationData,
+    NotificationMessageIds,
+    NotificationReponse,
+)
+from module.notification import Notifier
+from module.notification.base import NotificationContent
+from module.security.api import get_current_user
+
+router = APIRouter(
+    prefix="/notification",
+    tags=["notification"],
+    dependencies=[Depends(get_current_user)],
+)
+
+
+def get_notifier():
+    yield Notifier(
+        service_name=settings.notification.type,
+        config=settings.notification.dict(by_alias=True),
+    )
+
+
+@router.get("/total")
+async def get_total_notification(
+    notifier: Notifier = Depends(get_notifier),
+) -> NotificationReponse:
+    cursor = notifier.q.conn.cursor()
+    stmt = """SELECT COUNT(*) FROM Queue WHERE status=0"""
+
+    try:
+        total = cursor.execute(stmt).fetchone()[0]
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+    return dict(code=0, msg="success", data=dict(total=total))
+
+
+@router.get("")
+async def get_notification(
+    page: int = Query(1, ge=1),
+    limit: int = Query(20, ge=10, le=20, description="max limit is 20 per page"),
+    notifier: Notifier = Depends(get_notifier),
+) -> NotificationReponse:
+    cursor = notifier.q.conn.cursor()
+    stmt = r"""
+    SELECT 
+        message_id AS id, 
+        'AutoBangumi' AS title,
+        data AS content, 
+        in_time AS datetime, 
+        status AS has_read
+    FROM Queue
+    WHERE status=0
+    ORDER BY in_time DESC
+    """
+
+    offset = (page - 1) * limit
+    stmt += f"LIMIT {limit} OFFSET {offset}"
+
+    try:
+        rows = cursor.execute(stmt).fetchall()
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+    if not rows:
+        return NotificationReponse(
+            code=0, msg="success", data=dict(total=0, messages=[])
+        )
+
+    messages = [NotificationData(**data) for data in rows]
+
+    return NotificationReponse(
+        code=0, msg="success", data=dict(total=len(messages), messages=messages)
+    )
+
+
+@router.post("/read", response_model_exclude_none=True)
+async def set_notification_read(
+    body: NotificationMessageIds,
+    notifier: Notifier = Depends(get_notifier),
+) -> NotificationReponse:
+    try:
+        for message_id in body.message_ids:
+            notifier.q.done(message_id)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+    return NotificationReponse(code=0, msg="success")
+
+
+@router.post(
+    "/send",
+    description="send notification only for test",
+    response_model_exclude_none=True,
+)
+async def send_notification(
+    content: NotificationContent, notifier: Notifier = Depends(get_notifier)
+) -> NotificationReponse:
+    try:
+        await notifier.asend(content=content)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+    return NotificationReponse(code=0, msg="success")
+
+
+@router.get("/clean", response_model_exclude_none=True)
+async def clean_notification(
+    notifier: Notifier = Depends(get_notifier),
+) -> NotificationReponse:
+    try:
+        notifier.q.prune()
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+    return NotificationReponse(code=0, msg="success")
+
+
+# Note: put this vildcard route at the end of all routes to avoid masking other routes
+@router.get("/get")
+async def get_notification_by_id(
+    message_id: str, notifier: Notifier = Depends(get_notifier)
+) -> NotificationReponse:
+    message = notifier.q.get(message_id)
+    if not message:
+        raise HTTPException(status_code=404, detail="message not found")
+
+    return NotificationReponse(code=0, msg="success", data=message)

backend/src/module/conf/const.py

@@ -1,4 +1,5 @@
 # -*- encoding: utf-8 -*-
+import pathlib
 from urllib.parse import parse_qs, urlparse
 
 DEFAULT_SETTINGS = {
@@ -116,3 +117,6 @@ def _(color: str, *args: str) -> str:
     ENDC = "\033[0m"
     BOLD = "\033[1m"
     UNDERLINE = "\033[4m"
+
+
+ROOT = pathlib.Path(__file__).parent.parent.parent.absolute()

backend/src/module/core/sub_thread.py

@@ -1,10 +1,10 @@
 import threading
-import time
+from concurrent.futures import ThreadPoolExecutor
 
 from module.conf import settings
 from module.downloader import DownloadClient
 from module.manager import Renamer, eps_complete
-from module.notification import PostNotification
+from module.notification import Notifier
 from module.rss import RSSAnalyser, RSSEngine
 
 from .status import ProgramStatus
@@ -58,11 +58,17 @@ def rename_loop(self):
         while not self.stop_event.is_set():
             with Renamer() as renamer:
                 renamed_info = renamer.rename()
+
             if settings.notification.enable:
-                with PostNotification() as notifier:
-                    for info in renamed_info:
-                        notifier.send_msg(info)
-                        time.sleep(2)
+                notifier = Notifier(
+                    settings.notification.type,
+                    config=settings.notification.dict(by_alias=True),
+                )
+                with ThreadPoolExecutor(max_workers=2) as worker:
+                    worker.map(
+                        lambda info: notifier.send(notification=info), renamed_info
+                    )
+
             self.stop_event.wait(settings.program.rename_time)
 
     def rename_start(self):

backend/src/module/models/api.py

@@ -1,4 +1,7 @@
-from pydantic import BaseModel
+from datetime import datetime
+from typing import Any, Optional
+
+from pydantic import BaseModel, Field, validator
 
 
 class RssLink(BaseModel):
@@ -16,3 +19,27 @@ class ChangeConfig(BaseModel):
 
 class ChangeRule(BaseModel):
     rule: dict
+
+
+class NotificationMessageIds(BaseModel):
+    message_ids: list[str] = Field(..., description="message ids to be set read")
+
+
+class NotificationData(BaseModel):
+    title: str = Field("AutoBangumi", description="title of the notification")
+    content: str = Field(..., description="content of the notification")
+    datetime: str = Field(..., description="datetime of the notification")
+    has_read: bool = Field(False, description="whether the notification has been read")
+    id: str = Field(..., description="id of the notification")
+
+    @validator("datetime", pre=True)
+    def from_nanoseconds_to_datetime(cls, v: int) -> str:
+        return datetime.strftime(datetime.fromtimestamp(v / 1e9), "%Y-%m-%d %H:%M:%S")
+
+
+class NotificationReponse(BaseModel):
+    code: int = Field(200, description="status of the request")
+    msg: str = Field(..., description="message of the request")
+    data: dict[str, Any] = Field(
+        default_factory=dict, description="data of the request"
+    )

backend/src/module/models/config.py

@@ -3,6 +3,8 @@
 
 from pydantic import BaseModel, Field
 
+from module.notification.services import NotificationType
+
 
 class Program(BaseModel):
     rss_time: int = Field(900, description="Sleep time")
@@ -70,9 +72,11 @@ def password(self):
 
 class Notification(BaseModel):
     enable: bool = Field(False, description="Enable notification")
-    type: str = Field("telegram", description="Notification type")
+    type: NotificationType = Field("telegram", description="Notification type")
     token_: str = Field("", alias="token", description="Notification token")
     chat_id_: str = Field("", alias="chat_id", description="Notification chat id")
+    base_url: str = Field("", description="Notification base url")
+    channel: str = Field("", description="Notification channel")
 
     @property
     def token(self):