测试：编写消息通知功能

整体思路

其实做一个功能之前，还是最好规划下方案。因为对于我来说，这块能力还是相对薄弱了一些，所以有个大概的设计方案会有很大的帮助。

• Websocket

  由于我们需要服务器主动推送数据到客户端（浏览器），所以最好是服务器能主动通知，毕竟前端一直轮训接口，还是会有一定的性能损耗。

  如若我们在用户打开基础页面BasicLayout的时候，根据用户id获取websocket连接，并在服务端保持住。当有需要的时候则发送消息给用户。

  这里我们定义消息的类型: 广播消息和个人消息，广播消息类似系统通知那种，比如版本更新了什么，个人消息我目前能想到的就是测试计划执行了之后通知给对方，或者以后有关注case的功能，当case有变动，则通知到对应的人。

  所以我们需要一个存储所有连接的对象（字典），方便推送消息。

• 消息存储问题

  消息存储我们暂时还是先放到mysql，考虑到用户体量不会太多，我们可以只查询3个月的记录减缓数据的压力。

• 广播消息问题

  打算新开一个表存储用户阅读广播消息的记录，和消息表联表查出用户是否已读广播消息，如果用户选择查看全部消息的话，会比较复杂一些。

新增枚举文件app/enums/MessageEnum.py

from enum import IntEnum

class WebSocketMessageEnum(IntEnum):
    # 消息数量
    COUNT = 0
    # 桌面通知
    DESKTOP = 1

class MessageStateEnum(IntEnum):
    """
    消息状态枚举类
    """
    unread = 1  # 未读
    read = 2  # 已读

class MessageTypeEnum(IntEnum):
    """
    消息类型枚举类
    """
    all = 0  # 全部消息
    broadcast = 1  # 广播消息
    others = 2  # 其他消息

这边定义了3个枚举类

• 通知的内容类型

  消息有数量和桌面通知2种，由于可参考的资料比较少，加上我也没做过类似的项目，所以定义可能比较奇怪，大家仅供参考即可。

• 消息状态

  已读和未读。

• 消息类型

  全部消息/广播消息/其他消息（也可以说是个人消息）

编写消息返回体

app/core/msg/wss_msg.py

from app.enums.MessageEnum import WebSocketMessageEnum

class WebSocketMessage(object):

    @staticmethod
    def msg_count(count=1, total=False):
        return dict(type=WebSocketMessageEnum.COUNT, count=count, total=total)

    @staticmethod
    def desktop_msg(title, content=''):
        return dict(type=WebSocketMessageEnum.DESKTOP, title=title, content=content)

里面封装了2个方法:

• 消息数量

  因为我们有新消息来了，会告知对方来的消息数量，对方点入消息通知页面，可以看到所有数据。就好像知乎的邀请回答一样:

这边会有一个红色的数量上标，至于具体消息内容是啥，我们并不关心。用户需要点进去或者点开查看。

• 桌面通知

  桌面通知的话需要title和content，一个是标题，一个是正文。如果网站接收到对应的消息，则直接弹出桌面的对话框，区别就是这块是临时数据，不写入数据表。

  总的来说这2块的返回都是dict，在websocket里面我们会封装text，json和bytes3种消息格式的返回数据。

编写websocket管理类(由卫衣哥QYZHG编写，略有改动)

app/core/ws_connection_manager.py

# import abc
from typing import TypeVar

from fastapi import WebSocket

from app.core.msg.wss_msg import WebSocketMessage
from app.crud.notification.NotificationDao import PityNotificationDao
from app.models.notification import PityNotification
from app.utils.logger import Log

MsgType = TypeVar('MsgType', str, dict, bytes)

# class MsgSender(metaclass=abc.ABCMeta):
#     @abc.abstractmethod
#     def send_text(self):
#         pass
#
#     @abc.abstractmethod
#     def send_json(self):
#         pass
#
#     @abc.abstractmethod
#     def send_bytes(self):
#         pass

class ConnectionManager:
    BROADCAST = -1
    logger = Log("wss_manager")

    def __init__(self):
        self.active_connections: dict[int, WebSocket] = {}
        self.log = Log("websocket")

    async def connect(self, websocket: WebSocket, client_id: int) -> None:
        await websocket.accept()
        exist: WebSocket = self.active_connections.get(client_id)
        if exist:
            await exist.close()
            self.active_connections[client_id]: WebSocket = websocket
        else:
            self.active_connections[client_id]: WebSocket = websocket
            self.log.info(F"websocket:{client_id}： 建立连接成功！")

    def disconnect(self, client_id: int) -> None:
        del self.active_connections[client_id]
        self.log.info(F"websocket:{client_id}： 已安全断开！")

    @staticmethod
    async def pusher(sender: WebSocket, message: MsgType) -> None:
        """
        根据不同的消息类型，调用不同方法发送消息
        """
        msg_mapping: dict = {
            str: sender.send_text,
            dict: sender.send_json,
            bytes: sender.send_bytes
        }
        if func_push_msg := msg_mapping.get(type(message)):
            await func_push_msg(message)
        else:
            raise TypeError(F"websocket不能发送{type(message)}的内容！")

    async def send_personal_message(self, user_id: int, message: MsgType) -> None:
        """
        发送个人信息
        """
        conn = self.active_connections.get(user_id)
        if conn:
            await self.pusher(sender=conn, message=message)

    async def broadcast(self, message: MsgType) -> None:
        """
        广播
        """
        for connection in self.active_connections.values():
            await self.pusher(sender=connection, message=message)

    async def notify(self, user_id, title=None, content=None, notice: PityNotification = None):
        """
        根据user_id推送对应的
        :param content:
        :param title:
        :param user_id: 当user_id为-1的时候代表是广播消息
        :param notice:
        :return:
        """
        try:
            # 判断是否为桌面通知
            if title is not None:
                msg = WebSocketMessage.desktop_msg(title, content)
                if user_id == ConnectionManager.BROADCAST:
                    await self.broadcast(msg)
                else:
                    await self.send_personal_message(user_id, msg)
            else:
                # 说明不是桌面消息，直接给出消息数量即可
                if user_id == ConnectionManager.broadcast:
                    await self.broadcast(WebSocketMessage.msg_count())
                else:
                    await self.send_personal_message(user_id, WebSocketMessage.msg_count())
            # 判断是否要落入推送表
            if notice is not None:
                await PityNotificationDao.insert_record(notice)
        except Exception as e:
            ConnectionManager.logger.error(f"发送消息失败, {e}")


ws_manage = ConnectionManager()

里面的核心方法: notify是用来主动给网页发送消息的，稍后我们会说到。这里我们做了一个设定，如果title不为None，我们认定它是桌面通知类型，有可能有一些特殊的桌面通知也需要落库，所以我们再次判断，是否有Notification，传入了则说明需要入库。

当user_id为-1的时候，说明是一条广播消息。

调整消息表类

新增msg_title字段，msg_type为1时为广播消息(注释里面写的系统消息，广播消息更贴切一点)。

新增广播已读用户表

app/models/broadcast_read_user.py

from datetime import datetime

from sqlalchemy import Column, INT, DATETIME, BIGINT

from app.models import Base

class PityBroadcastReadUser(Base):
    id = Column(BIGINT, primary_key=True)
    notification_id = Column(INT, comment="对应消息id", index=True)
    read_user = Column(INT, comment="已读用户id")
    read_time = Column(DATETIME, comment="已读时间")

    __tablename__ = "pity_broadcast_read_user"

    def __init__(self, notification_id: int, read_user: int):
        self.notification_id = notification_id
        self.read_user = read_user
        self.read_time = datetime.now()
        self.id = None

这个表比较简单，只需要写入消息id+已读时间（其实都可以不要）+已读用户即可。

编写对应的dao类

app/crud/notification/BroadcastReadDao.py

from app.crud import Mapper
from app.models.broadcast_read_user import PityBroadcastReadUser
from app.utils.decorator import dao
from app.utils.logger import Log

@dao(PityBroadcastReadUser, Log("BroadcastReadDao"))
class BroadcastReadDao(Mapper):
    pass