Use error introduced in home-assistant#154579 in august integration (home-assistant#156096)

Author: wmoss

homeassistant/components/august/__init__.py

@@ -14,10 +14,11 @@
 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,
-    issue_registry as ir,
+from homeassistant.helpers import device_registry as dr, issue_registry as ir
+from homeassistant.helpers.config_entry_oauth2_flow import (
+    ImplementationUnavailableError,
+    OAuth2Session,
+    async_get_config_entry_implementation,
 )
 
 from .const import DEFAULT_AUGUST_BRAND, DOMAIN, PLATFORMS
@@ -37,14 +38,10 @@ async def async_setup_entry(hass: HomeAssistant, entry: AugustConfigEntry) -> bo
 
     session = async_create_august_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)
     august_gateway = AugustGateway(Path(hass.config.config_dir), session, oauth_session)
     try:
         await async_setup_august(hass, entry, august_gateway)

tests/components/august/mocks.py

@@ -121,9 +121,7 @@ def patch_august_setup():
     with (
         patch("yalexs.manager.gateway.ApiAsync") as api_mock,
         patch.object(_RateLimitChecker, "register_wakeup") as authenticate_mock,
-        patch(
-            "homeassistant.components.august.config_entry_oauth2_flow.async_get_config_entry_implementation"
-        ),
+        patch("homeassistant.components.august.async_get_config_entry_implementation"),
     ):
         yield api_mock, authenticate_mock
 

tests/components/august/test_init.py

@@ -24,6 +24,9 @@
     entity_registry as er,
     issue_registry as ir,
 )
+from homeassistant.helpers.config_entry_oauth2_flow import (
+    ImplementationUnavailableError,
+)
 from homeassistant.setup import async_setup_component
 
 from .mocks import (
@@ -294,8 +297,8 @@ async def test_oauth_implementation_not_available(hass: HomeAssistant) -> None:
     entry = await mock_august_config_entry(hass)
 
     with patch(
-        "homeassistant.components.august.config_entry_oauth2_flow.async_get_config_entry_implementation",
-        side_effect=ValueError("Implementation not available"),
+        "homeassistant.components.august.async_get_config_entry_implementation",
+        side_effect=ImplementationUnavailableError,
     ):
         await hass.config_entries.async_setup(entry.entry_id)
         await hass.async_block_till_done()