From 6fd9476bb9b0e18c11e44dfeccb9ee874362f173 Mon Sep 17 00:00:00 2001
From: jesperraemaekers <146726232+jesperraemaekers@users.noreply.github.com>
Date: Mon, 13 Jan 2025 14:01:57 +0100
Subject: [PATCH] Refresh token before setting up weheat (#135264)

---
 homeassistant/components/weheat/__init__.py | 17 ++++-
 tests/components/weheat/test_init.py        | 85 +++++++++++++++++++++
 2 files changed, 101 insertions(+), 1 deletion(-)
 create mode 100644 tests/components/weheat/test_init.py

diff --git a/homeassistant/components/weheat/__init__.py b/homeassistant/components/weheat/__init__.py
index a043a3a6845..37c1f721078 100644
--- a/homeassistant/components/weheat/__init__.py
+++ b/homeassistant/components/weheat/__init__.py
@@ -2,13 +2,16 @@
 
 from __future__ import annotations
 
+from http import HTTPStatus
+
+import aiohttp
 from weheat.abstractions.discovery import HeatPumpDiscovery
 from weheat.exceptions import UnauthorizedException
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_ACCESS_TOKEN, Platform
 from homeassistant.core import HomeAssistant
-from homeassistant.exceptions import ConfigEntryAuthFailed
+from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
 from homeassistant.helpers.config_entry_oauth2_flow import (
     OAuth2Session,
     async_get_config_entry_implementation,
@@ -28,6 +31,18 @@ async def async_setup_entry(hass: HomeAssistant, entry: WeheatConfigEntry) -> bo
 
     session = OAuth2Session(hass, entry, implementation)
 
+    try:
+        await session.async_ensure_token_valid()
+    except aiohttp.ClientResponseError as ex:
+        LOGGER.warning("API error: %s (%s)", ex.status, ex.message)
+        if ex.status in (
+            HTTPStatus.BAD_REQUEST,
+            HTTPStatus.UNAUTHORIZED,
+            HTTPStatus.FORBIDDEN,
+        ):
+            raise ConfigEntryAuthFailed("Token not valid, trigger renewal") from ex
+        raise ConfigEntryNotReady from ex
+
     token = session.token[CONF_ACCESS_TOKEN]
     entry.runtime_data = []
 
diff --git a/tests/components/weheat/test_init.py b/tests/components/weheat/test_init.py
new file mode 100644
index 00000000000..af5e2b8411b
--- /dev/null
+++ b/tests/components/weheat/test_init.py
@@ -0,0 +1,85 @@
+"""Tests for the weheat initialization."""
+
+from http import HTTPStatus
+from unittest.mock import AsyncMock, Mock, patch
+
+import pytest
+from weheat.abstractions.discovery import HeatPumpDiscovery
+
+from homeassistant.components.weheat import UnauthorizedException
+from homeassistant.config_entries import ConfigEntryState
+from homeassistant.core import HomeAssistant
+
+from . import setup_integration
+
+from tests.common import MockConfigEntry
+from tests.test_util.aiohttp import ClientResponseError
+
+
+@pytest.mark.usefixtures("setup_credentials")
+async def test_setup(
+    hass: HomeAssistant,
+    mock_weheat_discover: AsyncMock,
+    mock_weheat_heat_pump: AsyncMock,
+    mock_heat_pump_info: HeatPumpDiscovery.HeatPumpInfo,
+    mock_config_entry: MockConfigEntry,
+) -> None:
+    """Test the Weheat setup."""
+    mock_weheat_discover.return_value = [mock_heat_pump_info]
+
+    await setup_integration(hass, mock_config_entry)
+
+    assert mock_config_entry.state is ConfigEntryState.LOADED
+
+    await hass.config_entries.async_unload(mock_config_entry.entry_id)
+    await hass.async_block_till_done()
+    assert mock_config_entry.state is ConfigEntryState.NOT_LOADED
+
+
+@pytest.mark.usefixtures("setup_credentials")
+@pytest.mark.parametrize(
+    ("setup_exception", "expected_setup_state"),
+    [
+        (HTTPStatus.BAD_REQUEST, ConfigEntryState.SETUP_ERROR),
+        (HTTPStatus.UNAUTHORIZED, ConfigEntryState.SETUP_ERROR),
+        (HTTPStatus.FORBIDDEN, ConfigEntryState.SETUP_ERROR),
+        (HTTPStatus.GATEWAY_TIMEOUT, ConfigEntryState.SETUP_RETRY),
+    ],
+)
+async def test_setup_fail(
+    hass: HomeAssistant,
+    mock_weheat_discover: AsyncMock,
+    mock_weheat_heat_pump: AsyncMock,
+    mock_heat_pump_info: HeatPumpDiscovery.HeatPumpInfo,
+    mock_config_entry: MockConfigEntry,
+    setup_exception: Exception,
+    expected_setup_state: ConfigEntryState,
+) -> None:
+    """Test the Weheat setup with invalid token setup."""
+    with (
+        patch(
+            "homeassistant.components.weheat.OAuth2Session.async_ensure_token_valid",
+            side_effect=ClientResponseError(
+                Mock(real_url="http://example.com"), None, status=setup_exception
+            ),
+        ),
+    ):
+        await setup_integration(hass, mock_config_entry)
+
+    assert mock_config_entry.state is expected_setup_state
+
+
+@pytest.mark.usefixtures("setup_credentials")
+async def test_setup_fail_discover(
+    hass: HomeAssistant,
+    mock_weheat_discover: AsyncMock,
+    mock_weheat_heat_pump: AsyncMock,
+    mock_heat_pump_info: HeatPumpDiscovery.HeatPumpInfo,
+    mock_config_entry: MockConfigEntry,
+) -> None:
+    """Test the Weheat setup with and error from the heat pump discovery."""
+    mock_weheat_discover.side_effect = UnauthorizedException()
+
+    await setup_integration(hass, mock_config_entry)
+
+    assert mock_config_entry.state is ConfigEntryState.SETUP_ERROR