From 112d232c2e2240acc895dee8da4fa17c7b38c3f3 Mon Sep 17 00:00:00 2001
From: jjlawren <jjlawren@users.noreply.github.com>
Date: Tue, 29 Mar 2022 04:08:26 -0500
Subject: [PATCH] Autoresume in-progress items from Plex media browser (#68494)

---
 .../components/plex/media_browser.py          | 17 ++++++++++----
 homeassistant/components/plex/services.py     |  4 ++++
 tests/components/plex/fixtures/media_1.xml    |  2 +-
 tests/components/plex/test_browse_media.py    |  2 ++
 tests/components/plex/test_playback.py        | 16 +++++++++++++
 tests/components/plex/test_services.py        | 23 ++++++++++++++++++-
 6 files changed, 58 insertions(+), 6 deletions(-)

diff --git a/homeassistant/components/plex/media_browser.py b/homeassistant/components/plex/media_browser.py
index 4b2bb502d69..1df624e265b 100644
--- a/homeassistant/components/plex/media_browser.py
+++ b/homeassistant/components/plex/media_browser.py
@@ -66,7 +66,7 @@ def browse_media(  # noqa: C901
     if media_content_type in ("plex_root", None):
         return root_payload(hass, is_internal, platform=platform)
 
-    def item_payload(item, short_name=False):
+    def item_payload(item, short_name=False, extra_params=None):
         """Create response payload for a single media item."""
         try:
             media_class = ITEM_TYPE_MEDIA_CLASS[item.type]
@@ -75,7 +75,9 @@ def browse_media(  # noqa: C901
         payload = {
             "title": pretty_title(item, short_name),
             "media_class": media_class,
-            "media_content_id": generate_plex_uri(server_id, item.ratingKey),
+            "media_content_id": generate_plex_uri(
+                server_id, item.ratingKey, params=extra_params
+            ),
             "media_content_type": item.type,
             "can_play": True,
             "can_expand": item.type in EXPANDABLES,
@@ -209,7 +211,13 @@ def browse_media(  # noqa: C901
                     continue
                 payload["children"].append(station_payload(item))
             else:
-                payload["children"].append(item_payload(item))
+                extra_params = None
+                hub_context = hub.context.split(".")[-1]
+                if hub_context in ("continue", "inprogress", "ondeck"):
+                    extra_params = {"resume": 1}
+                payload["children"].append(
+                    item_payload(item, extra_params=extra_params)
+                )
         return BrowseMedia(**payload)
 
     if special_folder:
@@ -279,7 +287,7 @@ def browse_media(  # noqa: C901
     return response
 
 
-def generate_plex_uri(server_id, media_id):
+def generate_plex_uri(server_id, media_id, params=None):
     """Create a media_content_id URL for playable Plex media."""
     if isinstance(media_id, int):
         media_id = str(media_id)
@@ -290,6 +298,7 @@ def generate_plex_uri(server_id, media_id):
             scheme=DOMAIN,
             host=server_id,
             path=media_id,
+            query=params,
         )
     )
 
diff --git a/homeassistant/components/plex/services.py b/homeassistant/components/plex/services.py
index 5a6dd657942..8cf714b8823 100644
--- a/homeassistant/components/plex/services.py
+++ b/homeassistant/components/plex/services.py
@@ -112,6 +112,7 @@ def process_plex_payload(
 ) -> PlexMediaSearchResult:
     """Look up Plex media using media_player.play_media service payloads."""
     plex_server = default_plex_server
+    extra_params = {}
 
     if content_id.startswith(PLEX_URI_SCHEME + "{"):
         # Handle the special payload of 'plex://{<json>}'
@@ -132,6 +133,7 @@ def process_plex_payload(
         else:
             # Handle legacy payloads without server_id in URL host position
             content = int(plex_url.host)  # type: ignore[arg-type]
+        extra_params = dict(plex_url.query)
     else:
         content = json.loads(content_id)
 
@@ -152,6 +154,8 @@ def process_plex_payload(
         content = {"plex_key": content}
         content_type = DOMAIN
 
+    content.update(extra_params)
+
     if playqueue_id := content.pop("playqueue_id", None):
         if not supports_playqueues:
             raise HomeAssistantError("Plex playqueues are not supported on this device")
diff --git a/tests/components/plex/fixtures/media_1.xml b/tests/components/plex/fixtures/media_1.xml
index 838afb2959c..7392eaeef92 100644
--- a/tests/components/plex/fixtures/media_1.xml
+++ b/tests/components/plex/fixtures/media_1.xml
@@ -1,4 +1,4 @@
-<MediaContainer allowSync="1" identifier="com.plexapp.plugins.library" librarySectionID="1" librarySectionTitle="Movies" librarySectionUUID="805308ec-5019-43d4-a449-75d2b9e42f93" mediaTagPrefix="/system/bundle/media/flags/" mediaTagVersion="1603922053" size="1"><Video addedAt="1377829261" art="/library/metadata/1/art/1590245989" audienceRating="9.5" audienceRatingImage="rottentomatoes://image.rating.upright" chapterSource="agent" contentRating="R" duration="9000000" guid="com.plexapp.agents.imdb://tt0123456?lang=en" key="/library/metadata/1" lastViewedAt="1505969509" librarySectionID="1" librarySectionKey="/library/sections/1" librarySectionTitle="Movies" originallyAvailableAt="2000-01-01" primaryExtraKey="/library/metadata/195540" rating="9.0" ratingImage="rottentomatoes://image.rating.certified" ratingKey="1" studio="Studio Entertainment" summary="Some elaborate summary." tagline="Witty saying." thumb="/library/metadata/1/thumb/1590245989" title="Movie 1" type="movie" updatedAt="1590245989" viewCount="1" year="2000">
+<MediaContainer allowSync="1" identifier="com.plexapp.plugins.library" librarySectionID="1" librarySectionTitle="Movies" librarySectionUUID="805308ec-5019-43d4-a449-75d2b9e42f93" mediaTagPrefix="/system/bundle/media/flags/" mediaTagVersion="1603922053" size="1"><Video addedAt="1377829261" art="/library/metadata/1/art/1590245989" audienceRating="9.5" audienceRatingImage="rottentomatoes://image.rating.upright" chapterSource="agent" contentRating="R" duration="9000000" guid="com.plexapp.agents.imdb://tt0123456?lang=en" key="/library/metadata/1" lastViewedAt="1505969509" librarySectionID="1" librarySectionKey="/library/sections/1" librarySectionTitle="Movies" originallyAvailableAt="2000-01-01" primaryExtraKey="/library/metadata/195540" rating="9.0" ratingImage="rottentomatoes://image.rating.certified" ratingKey="1" studio="Studio Entertainment" summary="Some elaborate summary." tagline="Witty saying." thumb="/library/metadata/1/thumb/1590245989" title="Movie 1" type="movie" updatedAt="1590245989" viewCount="1" year="2000" viewOffset="555">
 <Genre count="119" filter="genre=25578" id="25578" tag="Sci-Fi" />
 <Genre count="197" filter="genre=87" id="87" tag="Action" />
 <Director count="4" filter="director=100" id="100" tag="Famous Director" />
diff --git a/tests/components/plex/test_browse_media.py b/tests/components/plex/test_browse_media.py
index 316c095160c..6af11a61191 100644
--- a/tests/components/plex/test_browse_media.py
+++ b/tests/components/plex/test_browse_media.py
@@ -204,6 +204,8 @@ async def test_browse_media(
     assert msg["success"]
     result = msg["result"]
     assert result[ATTR_MEDIA_CONTENT_TYPE] == "hub"
+    assert result["title"] == "Continue Watching"
+    assert result["children"][0]["media_content_id"].endswith("?resume=1")
 
     requests_mock.get(
         f"{mock_plex_server.url_in_use}/hubs/sections/3?includeStations=1",
diff --git a/tests/components/plex/test_playback.py b/tests/components/plex/test_playback.py
index e90a64a5fe1..10ab232f64f 100644
--- a/tests/components/plex/test_playback.py
+++ b/tests/components/plex/test_playback.py
@@ -127,6 +127,22 @@ async def test_media_player_playback(
     )
     assert playmedia_mock.called
 
+    # Test movie success with media browser URL and resuming
+    playmedia_mock.reset()
+    assert await hass.services.async_call(
+        MP_DOMAIN,
+        SERVICE_PLAY_MEDIA,
+        {
+            ATTR_ENTITY_ID: media_player,
+            ATTR_MEDIA_CONTENT_TYPE: MEDIA_TYPE_MOVIE,
+            ATTR_MEDIA_CONTENT_ID: PLEX_URI_SCHEME
+            + f"{DEFAULT_DATA[CONF_SERVER_IDENTIFIER]}/1?resume=1",
+        },
+        True,
+    )
+    assert playmedia_mock.called
+    assert playmedia_mock.last_request.qs["offset"][0] == "555"
+
     # Test movie success with legacy media browser URL
     playmedia_mock.reset()
     assert await hass.services.async_call(
diff --git a/tests/components/plex/test_services.py b/tests/components/plex/test_services.py
index f579fd90f90..f0224230c3a 100644
--- a/tests/components/plex/test_services.py
+++ b/tests/components/plex/test_services.py
@@ -21,7 +21,7 @@ from homeassistant.components.plex.services import process_plex_payload
 from homeassistant.const import CONF_URL
 from homeassistant.exceptions import HomeAssistantError
 
-from .const import DEFAULT_OPTIONS, SECONDARY_DATA
+from .const import DEFAULT_DATA, DEFAULT_OPTIONS, SECONDARY_DATA
 
 from tests.common import MockConfigEntry
 
@@ -210,3 +210,24 @@ async def test_lookup_media_for_other_integrations(
     requests_mock.post("/playqueues", text=playqueue_created)
     result = process_plex_payload(hass, MEDIA_TYPE_MUSIC, CONTENT_ID_SHUFFLE)
     assert isinstance(result.media, plexapi.playqueue.PlayQueue)
+
+
+async def test_lookup_media_with_urls(hass, mock_plex_server):
+    """Test media lookup for media_player.play_media calls from cast/sonos."""
+    CONTENT_ID_URL = f"{PLEX_URI_SCHEME}{DEFAULT_DATA[CONF_SERVER_IDENTIFIER]}/100"
+
+    # Test URL format
+    result = process_plex_payload(
+        hass, MEDIA_TYPE_MUSIC, CONTENT_ID_URL, supports_playqueues=False
+    )
+    assert isinstance(result.media, plexapi.audio.Track)
+    assert result.shuffle is False
+
+    # Test URL format with shuffle
+    CONTENT_ID_URL_WITH_SHUFFLE = CONTENT_ID_URL + "?shuffle=1"
+    result = process_plex_payload(
+        hass, MEDIA_TYPE_MUSIC, CONTENT_ID_URL_WITH_SHUFFLE, supports_playqueues=False
+    )
+    assert isinstance(result.media, plexapi.audio.Track)
+    assert result.shuffle is True
+    assert result.offset == 0