Add Plex library count sensors

homeassistant/components/plex/__init__.py

@@ -26,7 +26,10 @@
 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 @@
     PLATFORMS,
     PLATFORMS_COMPLETED,
     PLEX_SERVER_CONFIG,
+    PLEX_UPDATE_LIBRARY_SIGNAL,
     PLEX_UPDATE_PLATFORMS_SIGNAL,
     SERVERS,
     WEBSOCKETS,
@@ -179,12 +183,20 @@ def plex_websocket_callback(msgtype, data, error):
 
         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,
     )

homeassistant/components/plex/const.py

@@ -26,6 +26,7 @@
 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.{}"
 

homeassistant/components/plex/sensor.py

@@ -1,7 +1,10 @@
 """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
 
@@ -10,19 +13,45 @@
     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__)
 
 
 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):
@@ -103,6 +132,104 @@ def device_info(self):
             "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."""
+        unsub = async_dispatcher_connect(
+            self.hass,
+            PLEX_UPDATE_LIBRARY_SIGNAL.format(self.server_id),
+            self.async_refresh_sensor,
+        )
+        self.hass.data[PLEX_DOMAIN][DISPATCHERS][self.server_id].append(unsub)
+        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)
+        primary_libtype = LIBRARY_PRIMARY_LIBTYPE.get(
+            self.library_type, self.library_type
+        )
+
+        def query_library():
+            """Update library sensor state with sync calls."""
+            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
+                )
+
+        try:
+            await self.hass.async_add_executor_job(query_library)
+        except NotFound:
+            self._state = STATE_UNAVAILABLE
+        self.async_write_ha_state()
+
+    @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,
         }

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."""

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):

tests/components/plex/test_config_flow.py

@@ -650,6 +650,7 @@ def __init__(self):
         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"
 

tests/components/plex/test_sensor.py

@@ -0,0 +1,53 @@
+"""Tests for Plex sensors."""
+from homeassistant.const import STATE_UNAVAILABLE
+
+from .helpers import trigger_plex_update, wait_for_debouncer
+
+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"
+
+    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

tests/fixtures/plex/library_tvshows_size.xml

@@ -0,0 +1,3 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<MediaContainer size="0" totalSize="1" allowSync="1" art="/:/resources/show-fanart.jpg" identifier="com.plexapp.plugins.library" librarySectionID="2" librarySectionTitle="TV Shows" librarySectionUUID="905308ec-5019-43d4-a449-75d2b9e42f93" mediaTagPrefix="/system/bundle/media/flags/" mediaTagVersion="1614092584" nocache="1" offset="0" sortAsc="1" thumb="/:/resources/show.png" title1="TV Shows" title2="All Shows" viewGroup="show" viewMode="131122">
+</MediaContainer>

tests/fixtures/plex/library_tvshows_size_episodes.xml

@@ -0,0 +1,3 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<MediaContainer size="0" totalSize="10" allowSync="1" art="/:/resources/show-fanart.jpg" identifier="com.plexapp.plugins.library" librarySectionID="2" librarySectionTitle="TV Shows" librarySectionUUID="905308ec-5019-43d4-a449-75d2b9e42f93" mediaTagPrefix="/system/bundle/media/flags/" mediaTagVersion="1614092584" nocache="1" offset="0" sortAsc="1" thumb="/:/resources/show.png" title1="TV Shows" title2="All Shows" viewGroup="show" viewMode="131122">
+</MediaContainer>

tests/fixtures/plex/library_tvshows_size_seasons.xml

@@ -0,0 +1,3 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<MediaContainer size="0" totalSize="1" allowSync="1" art="/:/resources/show-fanart.jpg" identifier="com.plexapp.plugins.library" librarySectionID="2" librarySectionTitle="TV Shows" librarySectionUUID="905308ec-5019-43d4-a449-75d2b9e42f93" mediaTagPrefix="/system/bundle/media/flags/" mediaTagVersion="1614092584" nocache="1" offset="0" sortAsc="1" thumb="/:/resources/show.png" title1="TV Shows" title2="All Shows" viewGroup="show" viewMode="131122">
+</MediaContainer>