From 144cdf172c5bc404f46d774631e7ad43dea5c7fa Mon Sep 17 00:00:00 2001
From: qiangyanwen <508737091@qq.com>
Date: Tue, 13 Dec 2022 16:25:15 +0800
Subject: [PATCH] add test

---
 api/sys_monitor_api.py     | 40 ++++++++++++--------------------------
 utils/mem.py               | 14 ++-----------
 utils/websocket_manager.py |  7 ++++++-
 3 files changed, 20 insertions(+), 41 deletions(-)

diff --git a/api/sys_monitor_api.py b/api/sys_monitor_api.py
index 66d68fe..955ca93 100644
--- a/api/sys_monitor_api.py
+++ b/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 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
+        websocket: WebSocket
 ):
-    await websocket.accept()
-    while True:
-        data = await websocket.receive()
-        await websocket.send_text(f"消息是: {data}")
+    await manager.connect(websocket)
+    try:
+        while True:
+            data = await websocket.receive()
+            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)
diff --git a/utils/mem.py b/utils/mem.py
index 058b449..0eaca9d 100644
--- a/utils/mem.py
+++ b/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()
diff --git a/utils/websocket_manager.py b/utils/websocket_manager.py
index eb0237a..436ade1 100644
--- a/utils/websocket_manager.py
+++ b/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)