add test

api/sys_monitor_api.py

@@ -3,42 +3,26 @@
 # @Email :508737091@qq.com
 # @Author :qiangyanwen
 # @File :sys_monitor_api.py
+import json
 import time
 
 from fastapi import APIRouter, Body, Depends, Query, WebSocket, WebSocketDisconnect
 from utils.mem import get_disk_info, get_sys_info, get_cpu_info, ge_mem_info
-from utils.response import success_200
-from fastapi.responses import HTMLResponse
 from utils.websocket_manager import manager
 
 router = APIRouter(prefix="/api", tags=["系统模块"])
 
 
-@router.get("/sys/disk", summary="获取磁盘信息", name="获取磁盘信息")
-async def get_disk():
- return success_200(data=await get_disk_info(), message="查询磁盘信息成功")
-
-
-@router.get("/sys/info", summary="获取系统信息", name="获取系统信息")
-async def get_sys():
- return success_200(data=await get_sys_info(), message="查询系统信息成功")
-
-
-@router.get("/sys/cpu", summary="获取cpu信息", name="获取cpu信息")
-async def get_cpu():
- return success_200(data=await get_cpu_info(), message="查询cpu信息成功")
-
-
-@router.get("/sys/memory", summary="获取内存信息", name="获取内存信息")
-async def get_mem():
- return success_200(data=await ge_mem_info(), message="查询内存信息成功")
-
-
 @router.websocket("/ws")
 async def websocket_endpoint(
  websocket: WebSocket
 ):
- await websocket.accept()
+ await manager.connect(websocket)
+ try:
  while True:
  data = await websocket.receive()
- await websocket.send_text(f"消息是: {data}")
+ message = {"disk_info": await get_disk_info(), "sys_info": await get_sys_info(),
+ "cpu_info": await get_cpu_info(), "memory_info": await ge_mem_info()}
+ await manager.send_json(message, websocket)
+ except (WebSocketDisconnect, RuntimeError):
+ manager.disconnect(websocket)

utils/mem.py

@@ -14,7 +14,8 @@ async def get_sys_info():
  system_start_info = datetime.datetime.fromtimestamp(psutil.boot_time()).strftime("%Y-%m-%d %H:%M:%S")
  users_count = len(psutil.users())
  users_name = ",".join([u.name for u in psutil.users()])
- sys = SysModel(count=users_count, name=users_name, current_time=now_time, start_time=system_start_info)
+ # sys = SysModel(count=users_count, name=users_name, current_time=now_time, start_time=system_start_info)
+ sys = dict(count=users_count,name=users_name, current_time=now_time, start_time=system_start_info)
  return sys
 
 
@@ -65,14 +66,3 @@ async def get_disk_info():
  return disk_list
 
 
-if __name__ == '__main__':
- # now_time,users_info,system_start_info = get_sys_info()
- # print(now_time,users_info,system_start_info)
- cpu = get_cpu_info()
- print(cpu)
- # info = ge_mem_info()
- # print(info)
- # print(get_disk_info())
- # disk = get_disk_info()
- # print(disk)
- # get_net_info()

utils/websocket_manager.py

@@ -20,13 +20,18 @@ class WebsocketManager:
  await websocket.accept()
  self.active_connections.append(websocket)
 
- def disconnect(self, websocket: WebSocket):
+ async def disconnect(self, websocket: WebSocket):
+ await websocket.close()
  self.active_connections.remove(websocket)
 
  @staticmethod
  async def send_message(message: str, websocket: WebSocket):
  await websocket.send_text(message)
 
+ @staticmethod
+ async def send_json(message: str, websocket: WebSocket):
+ await websocket.send_json(message)
+
  async def broadcast(self, message: str):
  for connection in self.active_connections:
  await connection.send_text(message)