Coverage for custom_components/ha_repl_server/websocket_api.py: 54%

1"""Websocket commands: ha_repl_server/exec, ha_repl_server/reset, 

2ha_repl_server/sessions (all admin only).""" 

3 

4from __future__ import annotations 

5 

6from typing import Any, cast 

7 

8import probatio as vol 

9from homeassistant.components import websocket_api 

10from homeassistant.core import HomeAssistant, callback 

11 

12from .const import DEFAULT_SESSION, DOMAIN 

13from .session import SessionManager 

14 

15 

16@callback 

17def async_register(hass: HomeAssistant) -> None: 

18 websocket_api.async_register_command(hass, ws_exec) 

19 websocket_api.async_register_command(hass, ws_reset) 

20 websocket_api.async_register_command(hass, ws_sessions) 

21 

22 

23def _manager( 

24 hass: HomeAssistant, connection: websocket_api.ActiveConnection, msg: dict[str, Any] 

25) -> SessionManager | None: 

26 if (manager := hass.data.get(DOMAIN)) is None: 

27 connection.send_error( 

28 msg["id"], "not_loaded", "Home Assistant REPL integration is not loaded" 

29 ) 

30 return manager 

31 

32 

33@websocket_api.require_admin 

34# Home Assistant installs probatio as the `voluptuous` module itself 

35# (probatio.compat.install_as_voluptuous(), called from homeassistant's own 

36# __init__.py) - this schema and websocket_command()'s internal 

37# BASE_COMMAND_MESSAGE_SCHEMA.extend() both end up using the same probatio 

38# Marker classes at runtime, so this works. Static checkers don't know that: 

39# websocket_command()'s own signature is still annotated against real 

40# voluptuous's Marker/VolDictType, a type unrelated to probatio's from 

41# mypy/ty's point of view - hence the cast. 

42@websocket_api.websocket_command( 

43 cast( 

44 Any, 

45 { 

46 vol.Required("type"): "ha_repl_server/exec", 

47 vol.Required("code"): str, 

48 vol.Optional("session", default=DEFAULT_SESSION): str, 

49 vol.Optional("timeout"): vol.Coerce(float), 

50 vol.Optional("color", default=False): bool, 

51 vol.Optional("width", default=88): vol.Coerce(int), 

52 vol.Optional("auto_await", default=True): bool, 

53 }, 

54 ) 

55) 

56@websocket_api.async_response 

57async def ws_exec( 

58 hass: HomeAssistant, connection: websocket_api.ActiveConnection, msg: dict[str, Any] 

59) -> None: 

60 if (manager := _manager(hass, connection, msg)) is None: 

61 return 

62 result = await manager.get(msg["session"]).run( 

63 msg["code"], 

64 msg.get("timeout"), 

65 color=msg["color"], 

66 width=msg["width"], 

67 auto_await=msg["auto_await"], 

68 ) 

69 connection.send_result(msg["id"], result.as_dict()) 

70 

71 

72@websocket_api.require_admin 

73@websocket_api.websocket_command( 

74 cast( 

75 Any, 

76 { 

77 vol.Required("type"): "ha_repl_server/reset", 

78 vol.Optional("session", default=DEFAULT_SESSION): str, 

79 }, 

80 ) 

81) 

82@callback 

83def ws_reset( 

84 hass: HomeAssistant, connection: websocket_api.ActiveConnection, msg: dict[str, Any] 

85) -> None: 

86 if (manager := _manager(hass, connection, msg)) is None: 

87 return 

88 connection.send_result(msg["id"], {"reset": manager.reset(msg["session"])}) 

89 

90 

91@websocket_api.require_admin 

92@websocket_api.websocket_command( 

93 cast(Any, {vol.Required("type"): "ha_repl_server/sessions"}) 

94) 

95@callback 

96def ws_sessions( 

97 hass: HomeAssistant, connection: websocket_api.ActiveConnection, msg: dict[str, Any] 

98) -> None: 

99 if (manager := _manager(hass, connection, msg)) is None: 

100 return 

101 connection.send_result(msg["id"], {"sessions": manager.describe()})