def test_find_unserializable_data():
"""Find unserializeable data."""
assert find_paths_unserializable_data(1) == {}
assert find_paths_unserializable_data([1, 2]) == {}
assert find_paths_unserializable_data({"something": "yo"}) == {}
assert find_paths_unserializable_data({"something": set()}) == {
"$.something": set()
}
assert find_paths_unserializable_data({"something": [1, set()]}) == {
"$.something[1]": set()
}
assert find_paths_unserializable_data([1, {
"bla": set(),
"blub": set()
}]) == {
"$[1].bla": set(),
"$[1].blub": set(),
}
assert find_paths_unserializable_data({("A", ): 1}) == {
"$<key: ('A',)>": ("A", )
}
assert math.isnan(
find_paths_unserializable_data(float("nan"),
dump=partial(dumps,
allow_nan=False))["$"])
# Test custom encoder + State support.
class MockJSONEncoder(JSONEncoder):
"""Mock JSON encoder."""
def default(self, o):
"""Mock JSON encode method."""
if isinstance(o, datetime):
return o.isoformat()
return super().default(o)
bad_data = object()
assert (find_paths_unserializable_data(
[State("mock_domain.mock_entity", "on", {"bad": bad_data})],
dump=partial(dumps, cls=MockJSONEncoder),
) == {
"$[0](state: mock_domain.mock_entity).attributes.bad": bad_data
})
assert (find_paths_unserializable_data(
[Event("bad_event", {"bad_attribute": bad_data})],
dump=partial(dumps, cls=MockJSONEncoder),
) == {
"$[0](event: bad_event).data.bad_attribute": bad_data
})
def _get_json_file_response(
data: dict | list,
filename: str,
d_type: DiagnosticsType,
d_id: str,
sub_type: DiagnosticsSubType | None = None,
sub_id: str | None = None,
) -> web.Response:
"""Return JSON file from dictionary."""
try:
json_data = json.dumps(data, indent=2, cls=ExtendedJSONEncoder)
except TypeError:
_LOGGER.error(
"Failed to serialize to JSON: %s/%s%s. Bad data at %s",
d_type.value,
d_id,
f"/{sub_type.value}/{sub_id}" if sub_type is not None else "",
format_unserializable_data(find_paths_unserializable_data(data)),
)
return web.Response(status=HTTPStatus.INTERNAL_SERVER_ERROR)
return web.Response(
body=json_data,
content_type="application/json",
headers={"Content-Disposition": f'attachment; filename="{filename}.json"'},
)
def handle_get_states(hass: HomeAssistant, connection: ActiveConnection,
msg: dict[str, Any]) -> None:
"""Handle get states command."""
states = _async_get_allowed_states(hass, connection)
# JSON serialize here so we can recover if it blows up due to the
# state machine containing unserializable data. This command is required
# to succeed for the UI to show.
response = messages.result_message(msg["id"], states)
try:
connection.send_message(const.JSON_DUMP(response))
return
except (ValueError, TypeError):
connection.logger.error(
"Unable to serialize to JSON. Bad data found at %s",
format_unserializable_data(
find_paths_unserializable_data(response,
dump=const.JSON_DUMP)),
)
del response
# If we can't serialize, we'll filter out unserializable states
serialized = []
for state in states:
try:
serialized.append(const.JSON_DUMP(state))
except (ValueError, TypeError):
# Error is already logged above
pass
# We now have partially serialized states. Craft some JSON.
response2 = const.JSON_DUMP(
messages.result_message(msg["id"], ["TO_REPLACE"]))
response2 = response2.replace('"TO_REPLACE"', ", ".join(serialized))
connection.send_message(response2)
def handle_subscribe_entities(hass: HomeAssistant,
connection: ActiveConnection,
msg: dict[str, Any]) -> None:
"""Handle subscribe entities command."""
entity_ids = set(msg.get("entity_ids", []))
@callback
def forward_entity_changes(event: Event) -> None:
"""Forward entity state changed events to websocket."""
if not connection.user.permissions.check_entity(
event.data["entity_id"], POLICY_READ):
return
if entity_ids and event.data["entity_id"] not in entity_ids:
return
connection.send_message(
lambda: messages.cached_state_diff_message(msg["id"], event))
# We must never await between sending the states and listening for
# state changed events or we will introduce a race condition
# where some states are missed
states = _async_get_allowed_states(hass, connection)
connection.subscriptions[msg["id"]] = hass.bus.async_listen(
EVENT_STATE_CHANGED, forward_entity_changes, run_immediately=True)
connection.send_result(msg["id"])
data: dict[str, dict[str, dict]] = {
messages.ENTITY_EVENT_ADD: {
state.entity_id: messages.compressed_state_dict_add(state)
for state in states
if not entity_ids or state.entity_id in entity_ids
}
}
# JSON serialize here so we can recover if it blows up due to the
# state machine containing unserializable data. This command is required
# to succeed for the UI to show.
response = messages.event_message(msg["id"], data)
try:
connection.send_message(JSON_DUMP(response))
return
except (ValueError, TypeError):
connection.logger.error(
"Unable to serialize to JSON. Bad data found at %s",
format_unserializable_data(
find_paths_unserializable_data(response, dump=JSON_DUMP)),
)
del response
add_entities = data[messages.ENTITY_EVENT_ADD]
cannot_serialize: list[str] = []
for entity_id, state_dict in add_entities.items():
try:
JSON_DUMP(state_dict)
except (ValueError, TypeError):
cannot_serialize.append(entity_id)
for entity_id in cannot_serialize:
del add_entities[entity_id]
connection.send_message(JSON_DUMP(messages.event_message(msg["id"], data)))
async def _async_get_json_file_response(
hass: HomeAssistant,
data: dict | list,
filename: str,
domain: str,
d_type: DiagnosticsType,
d_id: str,
sub_type: DiagnosticsSubType | None = None,
sub_id: str | None = None,
) -> web.Response:
"""Return JSON file from dictionary."""
hass_sys_info = await async_get_system_info(hass)
hass_sys_info["run_as_root"] = hass_sys_info["user"] == "root"
del hass_sys_info["user"]
integration = await async_get_integration(hass, domain)
custom_components = {}
all_custom_components = await async_get_custom_components(hass)
for cc_domain, cc_obj in all_custom_components.items():
custom_components[cc_domain] = {
"version": cc_obj.version,
"requirements": cc_obj.requirements,
}
try:
json_data = json.dumps(
{
"home_assistant": hass_sys_info,
"custom_components": custom_components,
"integration_manifest": integration.manifest,
"data": data,
},
indent=2,
cls=ExtendedJSONEncoder,
)
except TypeError:
_LOGGER.error(
"Failed to serialize to JSON: %s/%s%s. Bad data at %s",
d_type.value,
d_id,
f"/{sub_type.value}/{sub_id}" if sub_type is not None else "",
format_unserializable_data(find_paths_unserializable_data(data)),
)
return web.Response(status=HTTPStatus.INTERNAL_SERVER_ERROR)
return web.Response(
body=json_data,
content_type="application/json",
headers={
"Content-Disposition":
f'attachment; filename="{filename}.json.txt"'
},
)
def message_to_json(message: Any) -> str:
"""Serialize a websocket message to json."""
try:
return const.JSON_DUMP(message)
except (ValueError, TypeError):
_LOGGER.error(
"Unable to serialize to JSON. Bad data found at %s",
format_unserializable_data(
find_paths_unserializable_data(message, dump=const.JSON_DUMP)),
)
return const.JSON_DUMP(
error_message(message["id"], const.ERR_UNKNOWN_ERROR,
"Invalid JSON in response"))
async def get( # pylint: disable=no-self-use
self, request: web.Request, d_type: str,
d_id: str) -> web.Response:
"""Download diagnostics."""
if d_type != "config_entry":
return web.Response(status=404)
hass = request.app["hass"]
config_entry = hass.config_entries.async_get_entry(d_id)
if config_entry is None:
return web.Response(status=404)
info = hass.data[DOMAIN].get(config_entry.domain)
if info is None:
return web.Response(status=404)
if info["config_entry"] is None:
return web.Response(status=404)
data = await info["config_entry"](hass, config_entry)
try:
json_data = json.dumps(data, indent=4, cls=ExtendedJSONEncoder)
except TypeError:
_LOGGER.error(
"Failed to serialize to JSON: %s/%s. Bad data at %s",
d_type,
d_id,
format_unserializable_data(
find_paths_unserializable_data(data)),
)
return web.Response(status=500)
return web.Response(
body=json_data,
content_type="application/json",
headers={
"Content-Disposition":
f'attachment; filename="{config_entry.domain}-{config_entry.entry_id}.json"'
},
)
def test_find_unserializable_data():
"""Find unserializeable data."""
assert find_paths_unserializable_data(1) == []
assert find_paths_unserializable_data([1, 2]) == []
assert find_paths_unserializable_data({"something": "yo"}) == []
assert find_paths_unserializable_data({"something":
set()}) == ["$.something"]
assert find_paths_unserializable_data({"something":
[1, set()]}) == ["$.something[1]"]
assert find_paths_unserializable_data([1, {
"bla": set(),
"blub": set()
}]) == [
"$[1].bla",
"$[1].blub",
]
assert find_paths_unserializable_data({("A", ): 1}) == ["$<key: ('A',)>"]
async def _writer(self):
"""Write outgoing messages."""
# Exceptions if Socket disconnected or cancelled by connection handler
with suppress(RuntimeError, ConnectionResetError, *CANCELLATION_ERRORS):
while not self.wsock.closed:
message = await self._to_write.get()
if message is None:
break
self._logger.debug("Sending %s", message)
if isinstance(message, str):
await self.wsock.send_str(message)
continue
try:
dumped = JSON_DUMP(message)
except (ValueError, TypeError):
await self.wsock.send_json(
error_message(
message["id"], ERR_UNKNOWN_ERROR, "Invalid JSON in response"
)
)
self._logger.error(
"Unable to serialize to JSON. Bad data found at %s",
format_unserializable_data(
find_paths_unserializable_data(message, dump=JSON_DUMP)
),
)
continue
await self.wsock.send_str(dumped)
# Clean up the peaker checker when we shut down the writer
if self._peak_checker_unsub:
self._peak_checker_unsub()
self._peak_checker_unsub = None
