From ebbe1ff1a29b02b356dcd94c39c2d9382ef473e4 Mon Sep 17 00:00:00 2001
From: Pedro Lamas <pedrolamas@gmail.com>
Date: Mon, 7 Feb 2022 15:49:18 +0000
Subject: [PATCH] Cache webostv supported_features state (#65930)

* Cache webostv supported_features state

* Fixes typings

* Restore supported_features attribute on restart

* Reverts change on supported_features initial state

Co-authored-by: Shay Levy <levyshay1@gmail.com>

* Fixes tests

Co-authored-by: Shay Levy <levyshay1@gmail.com>
---
 .../components/webostv/media_player.py        | 20 ++++++-
 tests/components/webostv/test_media_player.py | 54 +++++++++++++++++++
 2 files changed, 73 insertions(+), 1 deletion(-)

diff --git a/homeassistant/components/webostv/media_player.py b/homeassistant/components/webostv/media_player.py
index 5aac52f6f7b..8fa18ce3142 100644
--- a/homeassistant/components/webostv/media_player.py
+++ b/homeassistant/components/webostv/media_player.py
@@ -34,6 +34,7 @@ from homeassistant.components.media_player.const import (
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     ATTR_ENTITY_ID,
+    ATTR_SUPPORTED_FEATURES,
     ENTITY_MATCH_ALL,
     ENTITY_MATCH_NONE,
     STATE_OFF,
@@ -44,6 +45,7 @@ from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity import DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
+from homeassistant.helpers.restore_state import RestoreEntity
 
 from . import WebOsClientWrapper
 from .const import (
@@ -121,7 +123,7 @@ def cmd(
     return cmd_wrapper
 
 
-class LgWebOSMediaPlayerEntity(MediaPlayerEntity):
+class LgWebOSMediaPlayerEntity(RestoreEntity, MediaPlayerEntity):
     """Representation of a LG webOS Smart TV."""
 
     def __init__(
@@ -144,8 +146,12 @@ class LgWebOSMediaPlayerEntity(MediaPlayerEntity):
         self._current_source = None
         self._source_list: dict = {}
 
+        self._supported_features: int | None = None
+
     async def async_added_to_hass(self) -> None:
         """Connect and subscribe to dispatcher signals and state updates."""
+        await super().async_added_to_hass()
+
         self.async_on_remove(
             async_dispatcher_connect(self.hass, DOMAIN, self.async_signal_handler)
         )
@@ -154,6 +160,12 @@ class LgWebOSMediaPlayerEntity(MediaPlayerEntity):
             self.async_handle_state_update
         )
 
+        if self._supported_features is not None:
+            return
+
+        if (state := await self.async_get_last_state()) is not None:
+            self._supported_features = state.attributes.get(ATTR_SUPPORTED_FEATURES)
+
     async def async_will_remove_from_hass(self) -> None:
         """Call disconnect on removal."""
         self._client.unregister_state_update_callback(self.async_handle_state_update)
@@ -313,6 +325,9 @@ class LgWebOSMediaPlayerEntity(MediaPlayerEntity):
     @property
     def supported_features(self) -> int:
         """Flag media player features that are supported."""
+        if self.state == STATE_OFF and self._supported_features is not None:
+            return self._supported_features
+
         supported = SUPPORT_WEBOSTV
 
         if self._client.sound_output in ("external_arc", "external_speaker"):
@@ -323,6 +338,9 @@ class LgWebOSMediaPlayerEntity(MediaPlayerEntity):
         if self._wrapper.turn_on:
             supported |= SUPPORT_TURN_ON
 
+        if self.state != STATE_OFF:
+            self._supported_features = supported
+
         return supported
 
     @property
diff --git a/tests/components/webostv/test_media_player.py b/tests/components/webostv/test_media_player.py
index c249b491d9a..450d3d90377 100644
--- a/tests/components/webostv/test_media_player.py
+++ b/tests/components/webostv/test_media_player.py
@@ -556,3 +556,57 @@ async def test_supported_features(hass, client, monkeypatch):
     attrs = hass.states.get(ENTITY_ID).attributes
 
     assert attrs[ATTR_SUPPORTED_FEATURES] == supported
+
+
+async def test_cached_supported_features(hass, client, monkeypatch):
+    """Test test supported features."""
+    monkeypatch.setattr(client, "is_on", False)
+    monkeypatch.setattr(client, "sound_output", None)
+    await setup_webostv(hass)
+    await client.mock_state_update()
+
+    # TV off, support volume mute, step, set
+    supported = SUPPORT_WEBOSTV | SUPPORT_WEBOSTV_VOLUME | SUPPORT_VOLUME_SET
+    attrs = hass.states.get(ENTITY_ID).attributes
+
+    assert attrs[ATTR_SUPPORTED_FEATURES] == supported
+
+    # TV on, support volume mute, step
+    monkeypatch.setattr(client, "is_on", True)
+    monkeypatch.setattr(client, "sound_output", "external_speaker")
+    await client.mock_state_update()
+
+    supported = SUPPORT_WEBOSTV | SUPPORT_WEBOSTV_VOLUME
+    attrs = hass.states.get(ENTITY_ID).attributes
+
+    assert attrs[ATTR_SUPPORTED_FEATURES] == supported
+
+    # TV off, support volume mute, step
+    monkeypatch.setattr(client, "is_on", False)
+    monkeypatch.setattr(client, "sound_output", None)
+    await client.mock_state_update()
+
+    supported = SUPPORT_WEBOSTV | SUPPORT_WEBOSTV_VOLUME
+    attrs = hass.states.get(ENTITY_ID).attributes
+
+    assert attrs[ATTR_SUPPORTED_FEATURES] == supported
+
+    # TV on, support volume mute, step, set
+    monkeypatch.setattr(client, "is_on", True)
+    monkeypatch.setattr(client, "sound_output", "speaker")
+    await client.mock_state_update()
+
+    supported = SUPPORT_WEBOSTV | SUPPORT_WEBOSTV_VOLUME | SUPPORT_VOLUME_SET
+    attrs = hass.states.get(ENTITY_ID).attributes
+
+    assert attrs[ATTR_SUPPORTED_FEATURES] == supported
+
+    # TV off, support volume mute, step, step, set
+    monkeypatch.setattr(client, "is_on", False)
+    monkeypatch.setattr(client, "sound_output", None)
+    await client.mock_state_update()
+
+    supported = SUPPORT_WEBOSTV | SUPPORT_WEBOSTV_VOLUME | SUPPORT_VOLUME_SET
+    attrs = hass.states.get(ENTITY_ID).attributes
+
+    assert attrs[ATTR_SUPPORTED_FEATURES] == supported