From d62297a28b1309df8ed1e829b66b9a789c175920 Mon Sep 17 00:00:00 2001 From: jjlawren Date: Wed, 31 Mar 2021 06:57:16 -0500 Subject: [PATCH] Add Plex library count sensors (#48339) --- homeassistant/components/plex/__init__.py | 14 +- homeassistant/components/plex/const.py | 1 + homeassistant/components/plex/sensor.py | 151 ++++++++++++++++-- tests/components/plex/conftest.py | 18 +++ tests/components/plex/helpers.py | 4 +- tests/components/plex/test_config_flow.py | 1 + tests/components/plex/test_sensor.py | 76 +++++++++ tests/fixtures/plex/library_tvshows_size.xml | 3 + .../plex/library_tvshows_size_episodes.xml | 3 + .../plex/library_tvshows_size_seasons.xml | 3 + 10 files changed, 262 insertions(+), 12 deletions(-) create mode 100644 tests/components/plex/test_sensor.py create mode 100644 tests/fixtures/plex/library_tvshows_size.xml create mode 100644 tests/fixtures/plex/library_tvshows_size_episodes.xml create mode 100644 tests/fixtures/plex/library_tvshows_size_seasons.xml diff --git a/homeassistant/components/plex/__init__.py b/homeassistant/components/plex/__init__.py index 296f1594549..38825882fe9 100644 --- a/homeassistant/components/plex/__init__.py +++ b/homeassistant/components/plex/__init__.py @@ -26,7 +26,10 @@ from homeassistant.core import callback from homeassistant.exceptions import ConfigEntryNotReady from homeassistant.helpers.aiohttp_client import async_get_clientsession from homeassistant.helpers.debounce import Debouncer -from homeassistant.helpers.dispatcher import async_dispatcher_connect +from homeassistant.helpers.dispatcher import ( + async_dispatcher_connect, + async_dispatcher_send, +) from .const import ( CONF_SERVER, @@ -38,6 +41,7 @@ from .const import ( PLATFORMS, PLATFORMS_COMPLETED, PLEX_SERVER_CONFIG, + PLEX_UPDATE_LIBRARY_SIGNAL, PLEX_UPDATE_PLATFORMS_SIGNAL, SERVERS, WEBSOCKETS, @@ -179,12 +183,20 @@ async def async_setup_entry(hass, entry): elif msgtype == "playing": hass.async_create_task(plex_server.async_update_session(data)) + elif msgtype == "status": + if data["StatusNotification"][0]["title"] == "Library scan complete": + async_dispatcher_send( + hass, + PLEX_UPDATE_LIBRARY_SIGNAL.format(server_id), + ) session = async_get_clientsession(hass) + subscriptions = ["playing", "status"] verify_ssl = server_config.get(CONF_VERIFY_SSL) websocket = PlexWebsocket( plex_server.plex_server, plex_websocket_callback, + subscriptions=subscriptions, session=session, verify_ssl=verify_ssl, ) diff --git a/homeassistant/components/plex/const.py b/homeassistant/components/plex/const.py index eec433202e4..c8ab6c8f147 100644 --- a/homeassistant/components/plex/const.py +++ b/homeassistant/components/plex/const.py @@ -26,6 +26,7 @@ PLEX_SERVER_CONFIG = "server_config" PLEX_NEW_MP_SIGNAL = "plex_new_mp_signal.{}" PLEX_UPDATE_MEDIA_PLAYER_SESSION_SIGNAL = "plex_update_session_signal.{}" PLEX_UPDATE_MEDIA_PLAYER_SIGNAL = "plex_update_mp_signal.{}" +PLEX_UPDATE_LIBRARY_SIGNAL = "plex_update_libraries_signal.{}" PLEX_UPDATE_PLATFORMS_SIGNAL = "plex_update_platforms_signal.{}" PLEX_UPDATE_SENSOR_SIGNAL = "plex_update_sensor_signal.{}" diff --git a/homeassistant/components/plex/sensor.py b/homeassistant/components/plex/sensor.py index d79c4120dc7..c445548dc0a 100644 --- a/homeassistant/components/plex/sensor.py +++ b/homeassistant/components/plex/sensor.py @@ -1,19 +1,40 @@ """Support for Plex media server monitoring.""" import logging +from plexapi.exceptions import NotFound + from homeassistant.components.sensor import SensorEntity +from homeassistant.const import STATE_UNAVAILABLE from homeassistant.helpers.debounce import Debouncer from homeassistant.helpers.dispatcher import async_dispatcher_connect from .const import ( CONF_SERVER_IDENTIFIER, - DISPATCHERS, DOMAIN as PLEX_DOMAIN, NAME_FORMAT, + PLEX_UPDATE_LIBRARY_SIGNAL, PLEX_UPDATE_SENSOR_SIGNAL, SERVERS, ) +LIBRARY_ATTRIBUTE_TYPES = { + "artist": ["artist", "album"], + "photo": ["photoalbum"], + "show": ["show", "season"], +} + +LIBRARY_PRIMARY_LIBTYPE = { + "show": "episode", + "artist": "track", +} + +LIBRARY_ICON_LOOKUP = { + "artist": "mdi:music", + "movie": "mdi:movie", + "photo": "mdi:image", + "show": "mdi:television", +} + _LOGGER = logging.getLogger(__name__) @@ -21,8 +42,15 @@ async def async_setup_entry(hass, config_entry, async_add_entities): """Set up Plex sensor from a config entry.""" server_id = config_entry.data[CONF_SERVER_IDENTIFIER] plexserver = hass.data[PLEX_DOMAIN][SERVERS][server_id] - sensor = PlexSensor(hass, plexserver) - async_add_entities([sensor]) + sensors = [PlexSensor(hass, plexserver)] + + def create_library_sensors(): + """Create Plex library sensors with sync calls.""" + for library in plexserver.library.sections(): + sensors.append(PlexLibrarySectionSensor(hass, plexserver, library)) + + await hass.async_add_executor_job(create_library_sensors) + async_add_entities(sensors) class PlexSensor(SensorEntity): @@ -45,12 +73,13 @@ class PlexSensor(SensorEntity): async def async_added_to_hass(self): """Run when about to be added to hass.""" server_id = self._server.machine_identifier - unsub = async_dispatcher_connect( - self.hass, - PLEX_UPDATE_SENSOR_SIGNAL.format(server_id), - self.async_refresh_sensor, + self.async_on_remove( + async_dispatcher_connect( + self.hass, + PLEX_UPDATE_SENSOR_SIGNAL.format(server_id), + self.async_refresh_sensor, + ) ) - self.hass.data[PLEX_DOMAIN][DISPATCHERS][server_id].append(unsub) async def _async_refresh_sensor(self): """Set instance object and trigger an entity state update.""" @@ -103,6 +132,110 @@ class PlexSensor(SensorEntity): "identifiers": {(PLEX_DOMAIN, self._server.machine_identifier)}, "manufacturer": "Plex", "model": "Plex Media Server", - "name": "Activity Sensor", + "name": self._server.friendly_name, + "sw_version": self._server.version, + } + + +class PlexLibrarySectionSensor(SensorEntity): + """Representation of a Plex library section sensor.""" + + def __init__(self, hass, plex_server, plex_library_section): + """Initialize the sensor.""" + self._server = plex_server + self.server_name = plex_server.friendly_name + self.server_id = plex_server.machine_identifier + self.library_section = plex_library_section + self.library_type = plex_library_section.type + self._name = f"{self.server_name} Library - {plex_library_section.title}" + self._unique_id = f"library-{self.server_id}-{plex_library_section.uuid}" + self._state = None + self._attributes = {} + + async def async_added_to_hass(self): + """Run when about to be added to hass.""" + self.async_on_remove( + async_dispatcher_connect( + self.hass, + PLEX_UPDATE_LIBRARY_SIGNAL.format(self.server_id), + self.async_refresh_sensor, + ) + ) + await self.async_refresh_sensor() + + async def async_refresh_sensor(self): + """Update state and attributes for the library sensor.""" + _LOGGER.debug("Refreshing library sensor for '%s'", self.name) + try: + await self.hass.async_add_executor_job(self._update_state_and_attrs) + except NotFound: + self._state = STATE_UNAVAILABLE + self.async_write_ha_state() + + def _update_state_and_attrs(self): + """Update library sensor state with sync calls.""" + primary_libtype = LIBRARY_PRIMARY_LIBTYPE.get( + self.library_type, self.library_type + ) + + self._state = self.library_section.totalViewSize( + libtype=primary_libtype, includeCollections=False + ) + for libtype in LIBRARY_ATTRIBUTE_TYPES.get(self.library_type, []): + self._attributes[f"{libtype}s"] = self.library_section.totalViewSize( + libtype=libtype, includeCollections=False + ) + + @property + def entity_registry_enabled_default(self): + """Return if sensor should be enabled by default.""" + return False + + @property + def name(self): + """Return the name of the sensor.""" + return self._name + + @property + def unique_id(self): + """Return the id of this plex client.""" + return self._unique_id + + @property + def should_poll(self): + """Return True if entity has to be polled for state.""" + return False + + @property + def state(self): + """Return the state of the sensor.""" + return self._state + + @property + def unit_of_measurement(self): + """Return the unit this state is expressed in.""" + return "Items" + + @property + def icon(self): + """Return the icon of the sensor.""" + return LIBRARY_ICON_LOOKUP.get(self.library_type, "mdi:plex") + + @property + def extra_state_attributes(self): + """Return the state attributes.""" + return self._attributes + + @property + def device_info(self): + """Return a device description for device registry.""" + if self.unique_id is None: + return None + + return { + "identifiers": {(PLEX_DOMAIN, self.server_id)}, + "manufacturer": "Plex", + "model": "Plex Media Server", + "name": self.server_name, "sw_version": self._server.version, } diff --git a/tests/components/plex/conftest.py b/tests/components/plex/conftest.py index 295b560f4d1..358ab293bf0 100644 --- a/tests/components/plex/conftest.py +++ b/tests/components/plex/conftest.py @@ -120,6 +120,24 @@ def library_fixture(): return load_fixture("plex/library.xml") +@pytest.fixture(name="library_tvshows_size", scope="session") +def library_tvshows_size_fixture(): + """Load tvshow library size payload and return it.""" + return load_fixture("plex/library_tvshows_size.xml") + + +@pytest.fixture(name="library_tvshows_size_episodes", scope="session") +def library_tvshows_size_episodes_fixture(): + """Load tvshow library size in episodes payload and return it.""" + return load_fixture("plex/library_tvshows_size_episodes.xml") + + +@pytest.fixture(name="library_tvshows_size_seasons", scope="session") +def library_tvshows_size_seasons_fixture(): + """Load tvshow library size in seasons payload and return it.""" + return load_fixture("plex/library_tvshows_size_seasons.xml") + + @pytest.fixture(name="library_sections", scope="session") def library_sections_fixture(): """Load library sections payload and return it.""" diff --git a/tests/components/plex/helpers.py b/tests/components/plex/helpers.py index 35a01c3bfff..246922bccae 100644 --- a/tests/components/plex/helpers.py +++ b/tests/components/plex/helpers.py @@ -26,10 +26,10 @@ def websocket_connected(mock_websocket): callback(SIGNAL_CONNECTION_STATE, STATE_CONNECTED, None) -def trigger_plex_update(mock_websocket, payload=UPDATE_PAYLOAD): +def trigger_plex_update(mock_websocket, msgtype="playing", payload=UPDATE_PAYLOAD): """Call the websocket callback method with a Plex update.""" callback = mock_websocket.call_args[0][1] - callback("playing", payload, None) + callback(msgtype, payload, None) async def wait_for_debouncer(hass): diff --git a/tests/components/plex/test_config_flow.py b/tests/components/plex/test_config_flow.py index bdd78131800..e0555bab0e8 100644 --- a/tests/components/plex/test_config_flow.py +++ b/tests/components/plex/test_config_flow.py @@ -650,6 +650,7 @@ async def test_manual_config(hass, mock_plex_calls): result = await hass.config_entries.flow.async_configure( result["flow_id"], user_input=MANUAL_SERVER ) + await hass.async_block_till_done() assert result["type"] == "create_entry" diff --git a/tests/components/plex/test_sensor.py b/tests/components/plex/test_sensor.py new file mode 100644 index 00000000000..5fa50892f32 --- /dev/null +++ b/tests/components/plex/test_sensor.py @@ -0,0 +1,76 @@ +"""Tests for Plex sensors.""" +from datetime import timedelta + +from homeassistant.config_entries import RELOAD_AFTER_UPDATE_DELAY +from homeassistant.const import STATE_UNAVAILABLE +from homeassistant.helpers import entity_registry as er +from homeassistant.util import dt + +from .helpers import trigger_plex_update, wait_for_debouncer + +from tests.common import async_fire_time_changed + +LIBRARY_UPDATE_PAYLOAD = {"StatusNotification": [{"title": "Library scan complete"}]} + + +async def test_library_sensor_values( + hass, + setup_plex_server, + mock_websocket, + requests_mock, + library_tvshows_size, + library_tvshows_size_episodes, + library_tvshows_size_seasons, +): + """Test the library sensors.""" + requests_mock.get( + "/library/sections/2/all?includeCollections=0&type=2", + text=library_tvshows_size, + ) + requests_mock.get( + "/library/sections/2/all?includeCollections=0&type=3", + text=library_tvshows_size_seasons, + ) + requests_mock.get( + "/library/sections/2/all?includeCollections=0&type=4", + text=library_tvshows_size_episodes, + ) + + await setup_plex_server() + await wait_for_debouncer(hass) + + activity_sensor = hass.states.get("sensor.plex_plex_server_1") + assert activity_sensor.state == "1" + + # Ensure sensor is created as disabled + assert hass.states.get("sensor.plex_server_1_library_tv_shows") is None + + # Enable sensor and validate values + entity_registry = er.async_get(hass) + entity_registry.async_update_entity( + entity_id="sensor.plex_server_1_library_tv_shows", disabled_by=None + ) + await hass.async_block_till_done() + + async_fire_time_changed( + hass, + dt.utcnow() + timedelta(seconds=RELOAD_AFTER_UPDATE_DELAY + 1), + ) + await hass.async_block_till_done() + + library_tv_sensor = hass.states.get("sensor.plex_server_1_library_tv_shows") + assert library_tv_sensor.state == "10" + assert library_tv_sensor.attributes["seasons"] == 1 + assert library_tv_sensor.attributes["shows"] == 1 + + # Handle library deletion + requests_mock.get( + "/library/sections/2/all?includeCollections=0&type=2", status_code=404 + ) + trigger_plex_update( + mock_websocket, msgtype="status", payload=LIBRARY_UPDATE_PAYLOAD + ) + await hass.async_block_till_done() + + library_tv_sensor = hass.states.get("sensor.plex_server_1_library_tv_shows") + assert library_tv_sensor.state == STATE_UNAVAILABLE diff --git a/tests/fixtures/plex/library_tvshows_size.xml b/tests/fixtures/plex/library_tvshows_size.xml new file mode 100644 index 00000000000..49f3fb5ea9b --- /dev/null +++ b/tests/fixtures/plex/library_tvshows_size.xml @@ -0,0 +1,3 @@ + + + diff --git a/tests/fixtures/plex/library_tvshows_size_episodes.xml b/tests/fixtures/plex/library_tvshows_size_episodes.xml new file mode 100644 index 00000000000..b326ac0e721 --- /dev/null +++ b/tests/fixtures/plex/library_tvshows_size_episodes.xml @@ -0,0 +1,3 @@ + + + diff --git a/tests/fixtures/plex/library_tvshows_size_seasons.xml b/tests/fixtures/plex/library_tvshows_size_seasons.xml new file mode 100644 index 00000000000..49f3fb5ea9b --- /dev/null +++ b/tests/fixtures/plex/library_tvshows_size_seasons.xml @@ -0,0 +1,3 @@ + + +