From b44aafc2946ca93becaeab3214c232fc63124f76 Mon Sep 17 00:00:00 2001
From: Will Moss <wbmoss@gmail.com>
Date: Mon, 24 Nov 2025 22:28:04 -0800
Subject: [PATCH] Improved error handling for oauth2 configuration in neato
 integration (#156300)

---
 homeassistant/components/neato/__init__.py  | 20 +++++++----
 homeassistant/components/neato/strings.json |  5 +++
 tests/components/neato/test_init.py         | 40 +++++++++++++++++++++
 3 files changed, 58 insertions(+), 7 deletions(-)
 create mode 100644 tests/components/neato/test_init.py

diff --git a/homeassistant/components/neato/__init__.py b/homeassistant/components/neato/__init__.py
index e0a3f4bc37a..03d8dd5ab29 100644
--- a/homeassistant/components/neato/__init__.py
+++ b/homeassistant/components/neato/__init__.py
@@ -10,7 +10,11 @@ from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_TOKEN, Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
-from homeassistant.helpers import config_entry_oauth2_flow
+from homeassistant.helpers.config_entry_oauth2_flow import (
+    ImplementationUnavailableError,
+    OAuth2Session,
+    async_get_config_entry_implementation,
+)
 
 from . import api
 from .const import NEATO_DOMAIN, NEATO_LOGIN
@@ -33,13 +37,15 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     if CONF_TOKEN not in entry.data:
         raise ConfigEntryAuthFailed
 
-    implementation = (
-        await config_entry_oauth2_flow.async_get_config_entry_implementation(
-            hass, entry
-        )
-    )
+    try:
+        implementation = await async_get_config_entry_implementation(hass, entry)
+    except ImplementationUnavailableError as err:
+        raise ConfigEntryNotReady(
+            translation_domain=NEATO_DOMAIN,
+            translation_key="oauth2_implementation_unavailable",
+        ) from err
 
-    session = config_entry_oauth2_flow.OAuth2Session(hass, entry, implementation)
+    session = OAuth2Session(hass, entry, implementation)
     try:
         await session.async_ensure_token_valid()
     except aiohttp.ClientResponseError as ex:
diff --git a/homeassistant/components/neato/strings.json b/homeassistant/components/neato/strings.json
index 328f50b798b..6bd89650d01 100644
--- a/homeassistant/components/neato/strings.json
+++ b/homeassistant/components/neato/strings.json
@@ -46,6 +46,11 @@
       }
     }
   },
+  "exceptions": {
+    "oauth2_implementation_unavailable": {
+      "message": "[%key:common::exceptions::oauth2_implementation_unavailable::message%]"
+    }
+  },
   "services": {
     "custom_cleaning": {
       "description": "Starts a custom cleaning of your house.",
diff --git a/tests/components/neato/test_init.py b/tests/components/neato/test_init.py
new file mode 100644
index 00000000000..044ef1ced32
--- /dev/null
+++ b/tests/components/neato/test_init.py
@@ -0,0 +1,40 @@
+"""Tests for the Neato component."""
+
+from unittest.mock import patch
+
+from homeassistant.components.neato.const import NEATO_DOMAIN
+from homeassistant.config_entries import ConfigEntryState
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.config_entry_oauth2_flow import (
+    ImplementationUnavailableError,
+)
+
+from tests.common import MockConfigEntry
+
+
+async def test_oauth_implementation_not_available(
+    hass: HomeAssistant,
+) -> None:
+    """Test that unavailable OAuth implementation raises ConfigEntryNotReady."""
+    config_entry = MockConfigEntry(
+        domain=NEATO_DOMAIN,
+        data={
+            "auth_implementation": NEATO_DOMAIN,
+            "token": {
+                "refresh_token": "mock-refresh-token",
+                "access_token": "mock-access-token",
+                "type": "Bearer",
+                "expires_in": 60,
+            },
+        },
+    )
+    config_entry.add_to_hass(hass)
+
+    with patch(
+        "homeassistant.components.neato.async_get_config_entry_implementation",
+        side_effect=ImplementationUnavailableError,
+    ):
+        await hass.config_entries.async_setup(config_entry.entry_id)
+        await hass.async_block_till_done()
+
+    assert config_entry.state is ConfigEntryState.SETUP_RETRY