Use error introduced in home-assistant#154579 in yale integration (home-assistant#156095)

Author: wmoss

homeassistant/components/yale/__init__.py

@@ -16,7 +16,12 @@
 from homeassistant.const import EVENT_HOMEASSISTANT_STOP
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
-from homeassistant.helpers import config_entry_oauth2_flow, device_registry as dr
+from homeassistant.helpers import device_registry as dr
+from homeassistant.helpers.config_entry_oauth2_flow import (
+    ImplementationUnavailableError,
+    OAuth2Session,
+    async_get_config_entry_implementation,
+)
 
 from .const import DOMAIN, PLATFORMS
 from .data import YaleData
@@ -30,14 +35,10 @@ async def async_setup_entry(hass: HomeAssistant, entry: YaleConfigEntry) -> bool
     """Set up Yale from a config entry."""
     session = async_create_yale_clientsession(hass)
     try:
-        implementation = (
-            await config_entry_oauth2_flow.async_get_config_entry_implementation(
-                hass, entry
-            )
-        )
-    except ValueError as err:
+        implementation = await async_get_config_entry_implementation(hass, entry)
+    except ImplementationUnavailableError as err:
         raise ConfigEntryNotReady("OAuth implementation not available") from err
-    oauth_session = config_entry_oauth2_flow.OAuth2Session(hass, entry, implementation)
+    oauth_session = OAuth2Session(hass, entry, implementation)
     yale_gateway = YaleGateway(Path(hass.config.config_dir), session, oauth_session)
     try:
         await async_setup_yale(hass, entry, yale_gateway)

tests/components/yale/mocks.py

@@ -137,9 +137,7 @@ def patch_yale_setup():
         patch.object(_RateLimitChecker, "register_wakeup") as authenticate_mock,
         patch("yalexs.manager.data.SocketIORunner") as socketio_mock,
         patch.object(socketio_mock, "run"),
-        patch(
-            "homeassistant.components.yale.config_entry_oauth2_flow.async_get_config_entry_implementation"
-        ),
+        patch("homeassistant.components.yale.async_get_config_entry_implementation"),
     ):
         yield api_mock, authenticate_mock, socketio_mock
 

tests/components/yale/test_init.py

@@ -19,6 +19,9 @@
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers import device_registry as dr, entity_registry as er
+from homeassistant.helpers.config_entry_oauth2_flow import (
+    ImplementationUnavailableError,
+)
 from homeassistant.setup import async_setup_component
 
 from .mocks import (
@@ -244,8 +247,8 @@ async def test_oauth_implementation_not_available(hass: HomeAssistant) -> None:
     entry = await mock_yale_config_entry(hass)
 
     with patch(
-        "homeassistant.components.yale.config_entry_oauth2_flow.async_get_config_entry_implementation",
-        side_effect=ValueError("Implementation not available"),
+        "homeassistant.components.yale.async_get_config_entry_implementation",
+        side_effect=ImplementationUnavailableError,
     ):
         await hass.config_entries.async_setup(entry.entry_id)
         await hass.async_block_till_done()