Add entity_ids filter on get_states websocket command

homeassistant/components/websocket_api/commands.py

@@ -325,13 +325,23 @@ def _async_get_allowed_states(
 
 
 @callback
-@decorators.websocket_command({vol.Required("type"): "get_states"})
+@decorators.websocket_command(
+    {
+        vol.Required("type"): "get_states",
+        vol.Optional("entity_ids"): cv.entity_ids,
+    }
+)
 def handle_get_states(
     hass: HomeAssistant, connection: ActiveConnection, msg: dict[str, Any]
 ) -> None:
     """Handle get states command."""
     states = _async_get_allowed_states(hass, connection)
 
+    # filter states if entity_ids is set
+    entity_ids = set(msg.get("entity_ids", []))
+    if entity_ids:
+        states = [state for state in states if state.entity_id in entity_ids]
+
     try:
         serialized_states = [state.as_dict_json for state in states]
     except (ValueError, TypeError):

tests/components/websocket_api/test_commands.py

@@ -2916,3 +2916,25 @@ def auto_off_listener(event):
 
     await websocket_client.close()
     await hass.async_block_till_done()
+
+
+async def test_get_states_entity_id_filter(
+    hass: HomeAssistant, websocket_client: MockHAClientWebSocket
+) -> None:
+    """Test get_states command."""
+    hass.states.async_set("greeting.hello", "world")
+    hass.states.async_set("greeting.bye", "universe")
+
+    await websocket_client.send_json(
+        {
+            "id": 5,
+            "type": "get_states",
+            "entity_ids": ["greeting.bye", "greeting.missing"],
+        }
+    )
+
+    msg = await websocket_client.receive_json()
+    assert msg["id"] == 5
+    assert msg["type"] == const.TYPE_RESULT
+    assert msg["success"]
+    assert msg["result"] == [hass.states.get("greeting.bye").as_dict()]