Coverage for src/par_run/web.py: 71%

41 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-04-12 07:35 -0400

1"""Web UI Module""" 

2 

3from pathlib import Path 

4import rich 

5from fastapi import Body, FastAPI, Request, WebSocket 

6from fastapi.staticfiles import StaticFiles 

7from fastapi.templating import Jinja2Templates 

8 

9from .executor import Command, CommandGroup, ProcessingStrategy, read_commands_ini, write_commands_ini 

10 

11BASE_PATH = Path(__file__).resolve().parent 

12 

13ws_app = FastAPI() 

14ws_app.mount("/static", StaticFiles(directory=str(BASE_PATH / "static")), name="static") 

15templates = Jinja2Templates(directory=str(BASE_PATH / "templates")) 

16 

17 

18@ws_app.get("/") 

19async def ws_main(request: Request): 

20 """Get the main page.""" 

21 print("Main page") 

22 return templates.TemplateResponse("index.html", {"request": request}) 

23 

24 

25@ws_app.get("/get-commands-config") 

26async def get_commands_config(): 

27 """Get the commands configuration.""" 

28 return read_commands_ini("commands.ini") 

29 

30 

31@ws_app.post("/update-commands-config") 

32async def update_commands_config(updated_config: list[CommandGroup] = Body(...)): 

33 """Update the commands configuration.""" 

34 write_commands_ini("commands.ini", updated_config) 

35 return {"message": "Configuration updated successfully"} 

36 

37 

38class WebCommandCB: 

39 """Websocket command callbacks.""" 

40 

41 def __init__(self, ws: WebSocket): 

42 self.ws = ws 

43 

44 async def on_start(self, cmd: Command): 

45 rich.print(f"[blue bold]Started command {cmd.name}[/]") 

46 

47 async def on_recv(self, cmd: Command, output: str): 

48 await self.ws.send_json({"commandName": cmd.name, "output": output}) 

49 

50 async def on_term(self, cmd: Command, exit_code: int): 

51 await self.ws.send_json({"commandName": cmd.name, "output": {"ret_code": exit_code}}) 

52 

53 

54@ws_app.websocket_route("/ws") 

55async def websocket_endpoint(websocket: WebSocket): 

56 """Websocket endpoint to run commands.""" 

57 rich.print("Websocket connection") 

58 master_groups = read_commands_ini("commands.ini") 

59 print(master_groups) 

60 await websocket.accept() 

61 cb = WebCommandCB(websocket) 

62 print("Running commands") 

63 exit_code = 0 

64 for grp in master_groups: 

65 exit_code = exit_code or await grp.run_async(ProcessingStrategy.ON_RECV, cb)