diff --git a/homeassistant/components/websocket_api/messages.py b/homeassistant/components/websocket_api/messages.py index 52e97b60ccf..f68beff5924 100644 --- a/homeassistant/components/websocket_api/messages.py +++ b/homeassistant/components/websocket_api/messages.py @@ -12,6 +12,7 @@ from homeassistant.util.json import ( find_paths_unserializable_data, format_unserializable_data, ) +from homeassistant.util.yaml.loader import JSON_TYPE from . import const @@ -27,6 +28,9 @@ MINIMAL_MESSAGE_SCHEMA = vol.Schema( # Base schema to extend by message handlers BASE_COMMAND_MESSAGE_SCHEMA = vol.Schema({vol.Required("id"): cv.positive_int}) +IDEN_TEMPLATE = "__IDEN__" +IDEN_JSON_TEMPLATE = '"__IDEN__"' + def result_message(iden: int, result: Any = None) -> Dict: """Return a success result message.""" @@ -43,12 +47,11 @@ def error_message(iden: int, code: str, message: str) -> Dict: } -def event_message(iden: int, event: Any) -> Dict: +def event_message(iden: JSON_TYPE, event: Any) -> Dict: """Return an event message.""" return {"id": iden, "type": "event", "event": event} -@lru_cache(maxsize=128) def cached_event_message(iden: int, event: Event) -> str: """Return an event message. @@ -58,7 +61,17 @@ def cached_event_message(iden: int, event: Event) -> str: all getting many of the same events (mostly state changed) we can avoid serializing the same data for each connection. """ - return message_to_json(event_message(iden, event)) + return _cached_event_message(event).replace(IDEN_JSON_TEMPLATE, str(iden), 1) + + +@lru_cache(maxsize=128) +def _cached_event_message(event: Event) -> str: + """Cache and serialize the event to json. + + The IDEN_TEMPLATE is used which will be replaced + with the actual iden in cached_event_message + """ + return message_to_json(event_message(IDEN_TEMPLATE, event)) def message_to_json(message: Any) -> str: diff --git a/tests/components/websocket_api/test_messages.py b/tests/components/websocket_api/test_messages.py index 832b72c5c1c..3ec156e6949 100644 --- a/tests/components/websocket_api/test_messages.py +++ b/tests/components/websocket_api/test_messages.py @@ -1,6 +1,7 @@ """Test Websocket API messages module.""" from homeassistant.components.websocket_api.messages import ( + _cached_event_message as lru_event_cache, cached_event_message, message_to_json, ) @@ -24,6 +25,7 @@ async def test_cached_event_message(hass): await hass.async_block_till_done() assert len(events) == 2 + lru_event_cache.cache_clear() msg0 = cached_event_message(2, events[0]) assert msg0 == cached_event_message(2, events[0]) @@ -33,18 +35,49 @@ async def test_cached_event_message(hass): assert msg0 != msg1 - cache_info = cached_event_message.cache_info() + cache_info = lru_event_cache.cache_info() assert cache_info.hits == 2 assert cache_info.misses == 2 assert cache_info.currsize == 2 cached_event_message(2, events[1]) - cache_info = cached_event_message.cache_info() + cache_info = lru_event_cache.cache_info() assert cache_info.hits == 3 assert cache_info.misses == 2 assert cache_info.currsize == 2 +async def test_cached_event_message_with_different_idens(hass): + """Test that we cache event messages when the subscrition idens differ.""" + + events = [] + + @callback + def _event_listener(event): + events.append(event) + + hass.bus.async_listen(EVENT_STATE_CHANGED, _event_listener) + + hass.states.async_set("light.window", "on") + await hass.async_block_till_done() + + assert len(events) == 1 + + lru_event_cache.cache_clear() + + msg0 = cached_event_message(2, events[0]) + msg1 = cached_event_message(3, events[0]) + msg2 = cached_event_message(4, events[0]) + + assert msg0 != msg1 + assert msg0 != msg2 + + cache_info = lru_event_cache.cache_info() + assert cache_info.hits == 2 + assert cache_info.misses == 1 + assert cache_info.currsize == 1 + + async def test_message_to_json(caplog): """Test we can serialize websocket messages."""