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

1"""Web UI Module"""

2

3 from pathlib import Path

4

5 import rich

6 from fastapi import Body, FastAPI, Request, WebSocket

7 from fastapi.staticfiles import StaticFiles

8 from fastapi.templating import Jinja2Templates

9

10 from .executor import Command, ProcessingStrategy, read_commands_ini, write_commands_ini

11

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

13

14 ws_app = FastAPI()

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

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

17

18

19@ws_app.get("/")

20 async def ws_main(request: Request):

21 """Get the main page."""

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

23

24

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

26 async 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")

32 async def update_commands_config(updated_config=Body(...)): # noqa: B008

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

34 write_commands_ini("commands.ini", updated_config)

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

36

37

38 class 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")

55 async 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 await websocket.accept()

60 cb = WebCommandCB(websocket)

61 exit_code = 0

62 for grp in master_groups:

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