Improve Plex debounce/throttle logic

homeassistant/components/plex/server.py

@@ -1,5 +1,4 @@
 """Shared class to maintain Plex server instances."""
-from functools import partial, wraps
 import logging
 import ssl
 from urllib.parse import urlparse
@@ -13,8 +12,8 @@
 from homeassistant.components.media_player import DOMAIN as MP_DOMAIN
 from homeassistant.const import CONF_TOKEN, CONF_URL, CONF_VERIFY_SSL
 from homeassistant.core import callback
+from homeassistant.helpers.debounce import Debouncer
 from homeassistant.helpers.dispatcher import async_dispatcher_send
-from homeassistant.helpers.event import async_call_later
 
 from .const import (
     CONF_CLIENT_IDENTIFIER,
@@ -43,31 +42,6 @@
 plexapi.X_PLEX_VERSION = X_PLEX_VERSION
 
 
-def debounce(func):
-    """Decorate function to debounce callbacks from Plex websocket."""
-
-    unsub = None
-
-    async def call_later_listener(self, _):
-        """Handle call_later callback."""
-        nonlocal unsub
-        unsub = None
-        await func(self)
-
-    @wraps(func)
-    async def wrapper(self):
-        """Schedule async callback."""
-        nonlocal unsub
-        if unsub:
-            _LOGGER.debug("Throttling update of %s", self.friendly_name)
-            unsub()  # pylint: disable=not-callable
-        unsub = async_call_later(
-            self.hass, DEBOUNCE_TIMEOUT, partial(call_later_listener, self),
-        )
-
-    return wrapper
-
-
 class PlexServer:
     """Manages a single Plex server connection."""
 
@@ -87,6 +61,13 @@ def __init__(self, hass, server_config, known_server_id=None, options=None):
         self._accounts = []
         self._owner_username = None
         self._version = None
+        self.async_update_platforms = Debouncer(
+            hass,
+            _LOGGER,
+            cooldown=DEBOUNCE_TIMEOUT,
+            immediate=True,
+            function=self._async_update_platforms,
+        ).async_call
 
         # Header conditionally added as it is not available in config entry v1
         if CONF_CLIENT_IDENTIFIER in server_config:
@@ -192,8 +173,7 @@ def _fetch_platform_data(self):
         """Fetch all data from the Plex server in a single method."""
         return (self._plex_server.clients(), self._plex_server.sessions())
 
-    @debounce
-    async def async_update_platforms(self):
+    async def _async_update_platforms(self):
         """Update the platform entities."""
         _LOGGER.debug("Updating devices")
 

tests/components/plex/test_config_flow.py

@@ -15,13 +15,14 @@
     CONF_USE_EPISODE_ART,
     DOMAIN,
     PLEX_SERVER_CONFIG,
+    PLEX_UPDATE_PLATFORMS_SIGNAL,
     SERVERS,
 )
 from homeassistant.config_entries import ENTRY_STATE_LOADED
 from homeassistant.const import CONF_HOST, CONF_PORT, CONF_TOKEN, CONF_URL
+from homeassistant.helpers.dispatcher import async_dispatcher_send
 from homeassistant.setup import async_setup_component
 
-from .common import trigger_plex_update
 from .const import DEFAULT_DATA, DEFAULT_OPTIONS, MOCK_SERVERS, MOCK_TOKEN
 from .mock_classes import MockPlexAccount, MockPlexServer
 
@@ -415,7 +416,8 @@ async def test_option_flow_new_users_available(hass, caplog):
 
     server_id = mock_plex_server.machineIdentifier
 
-    await trigger_plex_update(hass, server_id)
+    async_dispatcher_send(hass, PLEX_UPDATE_PLATFORMS_SIGNAL.format(server_id))
+    await hass.async_block_till_done()
 
     monitored_users = hass.data[DOMAIN][SERVERS][server_id].option_monitored_users
 

tests/components/plex/test_init.py

@@ -3,8 +3,9 @@
 from datetime import timedelta
 import ssl
 
-from asynctest import patch
+from asynctest import ClockedTestCase, patch
 import plexapi
+import pytest
 import requests
 
 from homeassistant.components.media_player import DOMAIN as MP_DOMAIN
@@ -23,14 +24,19 @@
     CONF_URL,
     CONF_VERIFY_SSL,
 )
+from homeassistant.helpers.dispatcher import async_dispatcher_send
 from homeassistant.setup import async_setup_component
 import homeassistant.util.dt as dt_util
 
-from .common import trigger_plex_update
 from .const import DEFAULT_DATA, DEFAULT_OPTIONS, MOCK_SERVERS, MOCK_TOKEN
 from .mock_classes import MockPlexAccount, MockPlexServer
 
-from tests.common import MockConfigEntry, async_fire_time_changed
+from tests.common import (
+    MockConfigEntry,
+    async_fire_time_changed,
+    async_test_home_assistant,
+    mock_storage,
+)
 
 
 async def test_setup_with_config(hass):
@@ -67,70 +73,90 @@ async def test_setup_with_config(hass):
 
     assert loaded_server.plex_server == mock_plex_server
 
-    assert server_id in hass.data[const.DOMAIN][const.DISPATCHERS]
-    assert server_id in hass.data[const.DOMAIN][const.WEBSOCKETS]
-    assert (
-        hass.data[const.DOMAIN][const.PLATFORMS_COMPLETED][server_id] == const.PLATFORMS
-    )
 
+class TestClockedPlex(ClockedTestCase):
+    """Create clock-controlled asynctest class."""
 
-async def test_setup_with_config_entry(hass, caplog):
-    """Test setup component with config."""
+    @pytest.fixture(autouse=True)
+    def inject_fixture(self, caplog):
+        """Inject pytest fixtures as instance attributes."""
+        self.caplog = caplog
 
-    mock_plex_server = MockPlexServer()
+    async def setUp(self):
+        """Initialize this test class."""
+        self.hass = await async_test_home_assistant(self.loop)
+        self.mock_storage = mock_storage()
+        self.mock_storage.__enter__()
 
-    entry = MockConfigEntry(
-        domain=const.DOMAIN,
-        data=DEFAULT_DATA,
-        options=DEFAULT_OPTIONS,
-        unique_id=DEFAULT_DATA["server_id"],
-    )
+    async def tearDown(self):
+        """Clean up the HomeAssistant instance."""
+        await self.hass.async_stop()
+        self.mock_storage.__exit__(None, None, None)
 
-    with patch("plexapi.server.PlexServer", return_value=mock_plex_server), patch(
-        "homeassistant.components.plex.PlexWebsocket.listen"
-    ) as mock_listen:
-        entry.add_to_hass(hass)
-        assert await hass.config_entries.async_setup(entry.entry_id)
-        await hass.async_block_till_done()
+    async def test_setup_with_config_entry(self):
+        """Test setup component with config."""
+        hass = self.hass
 
-    assert mock_listen.called
+        mock_plex_server = MockPlexServer()
 
-    assert len(hass.config_entries.async_entries(const.DOMAIN)) == 1
-    assert entry.state == ENTRY_STATE_LOADED
+        entry = MockConfigEntry(
+            domain=const.DOMAIN,
+            data=DEFAULT_DATA,
+            options=DEFAULT_OPTIONS,
+            unique_id=DEFAULT_DATA["server_id"],
+        )
 
-    server_id = mock_plex_server.machineIdentifier
-    loaded_server = hass.data[const.DOMAIN][const.SERVERS][server_id]
+        with patch("plexapi.server.PlexServer", return_value=mock_plex_server), patch(
+            "homeassistant.components.plex.PlexWebsocket.listen"
+        ) as mock_listen:
+            entry.add_to_hass(hass)
+            assert await hass.config_entries.async_setup(entry.entry_id)
+            await hass.async_block_till_done()
 
-    assert loaded_server.plex_server == mock_plex_server
+        assert mock_listen.called
 
-    assert server_id in hass.data[const.DOMAIN][const.DISPATCHERS]
-    assert server_id in hass.data[const.DOMAIN][const.WEBSOCKETS]
-    assert (
-        hass.data[const.DOMAIN][const.PLATFORMS_COMPLETED][server_id] == const.PLATFORMS
-    )
+        assert len(hass.config_entries.async_entries(const.DOMAIN)) == 1
+        assert entry.state == ENTRY_STATE_LOADED
 
-    await trigger_plex_update(hass, server_id)
+        server_id = mock_plex_server.machineIdentifier
+        loaded_server = hass.data[const.DOMAIN][const.SERVERS][server_id]
 
-    sensor = hass.states.get("sensor.plex_plex_server_1")
-    assert sensor.state == str(len(mock_plex_server.accounts))
+        assert loaded_server.plex_server == mock_plex_server
+
+        async_dispatcher_send(
+            hass, const.PLEX_UPDATE_PLATFORMS_SIGNAL.format(server_id)
+        )
+        await hass.async_block_till_done()
 
-    await trigger_plex_update(hass, server_id)
+        sensor = hass.states.get("sensor.plex_plex_server_1")
+        assert sensor.state == str(len(mock_plex_server.accounts))
 
-    for test_exception in (
-        plexapi.exceptions.BadRequest,
-        requests.exceptions.RequestException,
-    ):
-        with patch.object(
-            mock_plex_server, "clients", side_effect=test_exception
-        ) as patched_clients_bad_request:
-            await trigger_plex_update(hass, server_id)
-
-        assert patched_clients_bad_request.called
-        assert (
-            f"Could not connect to Plex server: {mock_plex_server.friendlyName}"
-            in caplog.text
+        # Ensure existing entities refresh
+        await self.advance(const.DEBOUNCE_TIMEOUT)
+        async_dispatcher_send(
+            hass, const.PLEX_UPDATE_PLATFORMS_SIGNAL.format(server_id)
         )
-        caplog.clear()
+        await hass.async_block_till_done()
+
+        for test_exception in (
+            plexapi.exceptions.BadRequest,
+            requests.exceptions.RequestException,
+        ):
+            with patch.object(
+                mock_plex_server, "clients", side_effect=test_exception
+            ) as patched_clients_bad_request:
+                await self.advance(const.DEBOUNCE_TIMEOUT)
+                async_dispatcher_send(
+                    hass, const.PLEX_UPDATE_PLATFORMS_SIGNAL.format(server_id)
+                )
+                await hass.async_block_till_done()
+
+            assert patched_clients_bad_request.called
+            assert (
+                f"Could not connect to Plex server: {mock_plex_server.friendlyName}"
+                in self.caplog.text
+            )
+            self.caplog.clear()
 
 
 async def test_set_config_entry_unique_id(hass):
@@ -251,22 +277,12 @@ async def test_unload_config_entry(hass):
 
     assert loaded_server.plex_server == mock_plex_server
 
-    assert server_id in hass.data[const.DOMAIN][const.DISPATCHERS]
-    assert server_id in hass.data[const.DOMAIN][const.WEBSOCKETS]
-    assert (
-        hass.data[const.DOMAIN][const.PLATFORMS_COMPLETED][server_id] == const.PLATFORMS
-    )
-
     with patch("homeassistant.components.plex.PlexWebsocket.close") as mock_close:
         await hass.config_entries.async_unload(entry.entry_id)
         assert mock_close.called
 
     assert entry.state == ENTRY_STATE_NOT_LOADED
 
-    assert server_id not in hass.data[const.DOMAIN][const.SERVERS]
-    assert server_id not in hass.data[const.DOMAIN][const.DISPATCHERS]
-    assert server_id not in hass.data[const.DOMAIN][const.WEBSOCKETS]
-
 
 async def test_setup_with_photo_session(hass):
     """Test setup component with config."""
@@ -292,7 +308,8 @@ async def test_setup_with_photo_session(hass):
 
     server_id = mock_plex_server.machineIdentifier
 
-    await trigger_plex_update(hass, server_id)
+    async_dispatcher_send(hass, const.PLEX_UPDATE_PLATFORMS_SIGNAL.format(server_id))
+    await hass.async_block_till_done()
 
     media_player = hass.states.get("media_player.plex_product_title")
     assert media_player.state == "idle"