From c1213857ced54041b883b8b8134ae1f58dea5273 Mon Sep 17 00:00:00 2001
From: jjlawren <jjlawren@users.noreply.github.com>
Date: Tue, 18 Oct 2022 04:51:43 -0500
Subject: [PATCH] Fix Plex reauth with multiple available servers (#80508)

---
 homeassistant/components/plex/config_flow.py |  8 ++-
 tests/components/plex/test_config_flow.py    | 56 ++++++++++++++++++++
 2 files changed, 63 insertions(+), 1 deletion(-)

diff --git a/homeassistant/components/plex/config_flow.py b/homeassistant/components/plex/config_flow.py
index e79b7e7ee04..1ebe439ff7c 100644
--- a/homeassistant/components/plex/config_flow.py
+++ b/homeassistant/components/plex/config_flow.py
@@ -106,6 +106,7 @@ class PlexFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
         self.token = None
         self.client_id = None
         self._manual = False
+        self._reauth_config = None
 
     async def async_step_user(self, user_input=None, errors=None):
         """Handle a flow initialized by the user."""
@@ -178,6 +179,9 @@ class PlexFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
 
     async def async_step_server_validate(self, server_config):
         """Validate a provided configuration."""
+        if self._reauth_config:
+            server_config = {**self._reauth_config, **server_config}
+
         errors = {}
         self.current_login = server_config
 
@@ -336,7 +340,9 @@ class PlexFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
 
     async def async_step_reauth(self, entry_data: Mapping[str, Any]) -> FlowResult:
         """Handle a reauthorization flow request."""
-        self.current_login = dict(entry_data)
+        self._reauth_config = {
+            CONF_SERVER_IDENTIFIER: entry_data[CONF_SERVER_IDENTIFIER]
+        }
         return await self.async_step_user()
 
 
diff --git a/tests/components/plex/test_config_flow.py b/tests/components/plex/test_config_flow.py
index fb5a0f06724..8c25baa8746 100644
--- a/tests/components/plex/test_config_flow.py
+++ b/tests/components/plex/test_config_flow.py
@@ -766,6 +766,62 @@ async def test_trigger_reauth(
     assert entry.data[PLEX_SERVER_CONFIG][CONF_TOKEN] == "BRAND_NEW_TOKEN"
 
 
+async def test_trigger_reauth_multiple_servers_available(
+    hass,
+    entry,
+    mock_plex_server,
+    mock_websocket,
+    current_request_with_host,
+    requests_mock,
+    plextv_resources_two_servers,
+):
+    """Test setup and reauthorization of a Plex token when multiple servers are available."""
+    assert entry.state is ConfigEntryState.LOADED
+
+    requests_mock.get(
+        "https://plex.tv/api/resources",
+        text=plextv_resources_two_servers,
+    )
+
+    with patch(
+        "plexapi.server.PlexServer.clients", side_effect=plexapi.exceptions.Unauthorized
+    ), patch("plexapi.server.PlexServer", side_effect=plexapi.exceptions.Unauthorized):
+        trigger_plex_update(mock_websocket)
+        await wait_for_debouncer(hass)
+
+    assert len(hass.config_entries.async_entries(DOMAIN)) == 1
+    assert entry.state is not ConfigEntryState.LOADED
+
+    flows = hass.config_entries.flow.async_progress()
+    assert len(flows) == 1
+    assert flows[0]["context"]["source"] == SOURCE_REAUTH
+
+    flow_id = flows[0]["flow_id"]
+
+    with patch("plexauth.PlexAuth.initiate_auth"), patch(
+        "plexauth.PlexAuth.token", return_value="BRAND_NEW_TOKEN"
+    ):
+        result = await hass.config_entries.flow.async_configure(flow_id, user_input={})
+        assert result["type"] == "external"
+
+        result = await hass.config_entries.flow.async_configure(result["flow_id"])
+        assert result["type"] == "external_done"
+
+        result = await hass.config_entries.flow.async_configure(result["flow_id"])
+        assert result["type"] == "abort"
+        assert result["flow_id"] == flow_id
+        assert result["reason"] == "reauth_successful"
+
+    assert len(hass.config_entries.flow.async_progress()) == 0
+    assert len(hass.config_entries.async_entries(DOMAIN)) == 1
+
+    assert entry.state is ConfigEntryState.LOADED
+    assert entry.data[CONF_SERVER] == mock_plex_server.friendly_name
+    assert entry.data[CONF_SERVER_IDENTIFIER] == mock_plex_server.machine_identifier
+    assert entry.data[PLEX_SERVER_CONFIG][CONF_URL] == PLEX_DIRECT_URL
+    assert entry.data[PLEX_SERVER_CONFIG][CONF_TOKEN] == "BRAND_NEW_TOKEN"
+
+
 async def test_client_request_missing(hass):
     """Test when client headers are not set properly."""
     result = await hass.config_entries.flow.async_init(