diff --git a/.coveragerc b/.coveragerc index d8330e13fdf..6c40a3e043d 100644 --- a/.coveragerc +++ b/.coveragerc @@ -873,6 +873,7 @@ omit = homeassistant/components/notion/__init__.py homeassistant/components/notion/binary_sensor.py homeassistant/components/notion/sensor.py + homeassistant/components/notion/util.py homeassistant/components/nsw_fuel_station/sensor.py homeassistant/components/nuki/__init__.py homeassistant/components/nuki/binary_sensor.py diff --git a/homeassistant/components/notion/__init__.py b/homeassistant/components/notion/__init__.py index bd83b192a69..3ed2c7bdb93 100644 --- a/homeassistant/components/notion/__init__.py +++ b/homeassistant/components/notion/__init__.py @@ -7,7 +7,6 @@ from datetime import timedelta from typing import Any from uuid import UUID -from aionotion import async_get_client from aionotion.bridge.models import Bridge from aionotion.errors import InvalidCredentialsError, NotionError from aionotion.listener.models import Listener, ListenerKind @@ -19,7 +18,6 @@ from homeassistant.const import CONF_PASSWORD, CONF_USERNAME, Platform from homeassistant.core import HomeAssistant, callback from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady from homeassistant.helpers import ( - aiohttp_client, config_validation as cv, device_registry as dr, entity_registry as er, @@ -33,6 +31,8 @@ from homeassistant.helpers.update_coordinator import ( ) from .const import ( + CONF_REFRESH_TOKEN, + CONF_USER_UUID, DOMAIN, LOGGER, SENSOR_BATTERY, @@ -46,6 +46,7 @@ from .const import ( SENSOR_TEMPERATURE, SENSOR_WINDOW_HINGED, ) +from .util import async_get_client_with_credentials, async_get_client_with_refresh_token PLATFORMS = [Platform.BINARY_SENSOR, Platform.SENSOR] @@ -139,25 +140,48 @@ class NotionData: async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: """Set up Notion as a config entry.""" - if not entry.unique_id: - hass.config_entries.async_update_entry( - entry, unique_id=entry.data[CONF_USERNAME] - ) + entry_updates: dict[str, Any] = {"data": {**entry.data}} - session = aiohttp_client.async_get_clientsession(hass) + if not entry.unique_id: + entry_updates["unique_id"] = entry.data[CONF_USERNAME] try: - client = await async_get_client( - entry.data[CONF_USERNAME], - entry.data[CONF_PASSWORD], - session=session, - use_legacy_auth=True, - ) + if password := entry_updates["data"].pop(CONF_PASSWORD, None): + # If a password exists in the config entry data, use it to get a new client + # (and pop it from the new entry data): + client = await async_get_client_with_credentials( + hass, entry.data[CONF_USERNAME], password + ) + else: + # If a password doesn't exist in the config entry data, we can safely assume + # that a refresh token and user UUID do, so we use them to get the client: + client = await async_get_client_with_refresh_token( + hass, + entry.data[CONF_USER_UUID], + entry.data[CONF_REFRESH_TOKEN], + ) except InvalidCredentialsError as err: - raise ConfigEntryAuthFailed("Invalid username and/or password") from err + raise ConfigEntryAuthFailed("Invalid credentials") from err except NotionError as err: raise ConfigEntryNotReady("Config entry failed to load") from err + # Always update the config entry with the latest refresh token and user UUID: + entry_updates["data"][CONF_REFRESH_TOKEN] = client.refresh_token + entry_updates["data"][CONF_USER_UUID] = client.user_uuid + + @callback + def async_save_refresh_token(refresh_token: str) -> None: + """Save a refresh token to the config entry data.""" + LOGGER.debug("Saving new refresh token to HASS storage") + hass.config_entries.async_update_entry( + entry, data={**entry.data, CONF_REFRESH_TOKEN: refresh_token} + ) + + # Create a callback to save the refresh token when it changes: + entry.async_on_unload(client.add_refresh_token_callback(async_save_refresh_token)) + + hass.config_entries.async_update_entry(entry, **entry_updates) + async def async_update() -> NotionData: """Get the latest data from the Notion API.""" data = NotionData(hass=hass, entry=entry) diff --git a/homeassistant/components/notion/config_flow.py b/homeassistant/components/notion/config_flow.py index 2ed83adeb08..f43c87b5085 100644 --- a/homeassistant/components/notion/config_flow.py +++ b/homeassistant/components/notion/config_flow.py @@ -2,9 +2,9 @@ from __future__ import annotations from collections.abc import Mapping +from dataclasses import dataclass, field from typing import Any -from aionotion import async_get_client from aionotion.errors import InvalidCredentialsError, NotionError import voluptuous as vol @@ -13,9 +13,9 @@ from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_PASSWORD, CONF_USERNAME from homeassistant.core import HomeAssistant from homeassistant.data_entry_flow import FlowResult -from homeassistant.helpers import aiohttp_client -from .const import DOMAIN, LOGGER +from .const import CONF_REFRESH_TOKEN, CONF_USER_UUID, DOMAIN, LOGGER +from .util import async_get_client_with_credentials AUTH_SCHEMA = vol.Schema( { @@ -30,17 +30,23 @@ REAUTH_SCHEMA = vol.Schema( ) +@dataclass(frozen=True, kw_only=True) +class CredentialsValidationResult: + """Define a validation result.""" + + user_uuid: str | None = None + refresh_token: str | None = None + errors: dict[str, Any] = field(default_factory=dict) + + async def async_validate_credentials( hass: HomeAssistant, username: str, password: str -) -> dict[str, Any]: - """Validate a Notion username and password (returning any errors).""" - session = aiohttp_client.async_get_clientsession(hass) +) -> CredentialsValidationResult: + """Validate a Notion username and password.""" errors = {} try: - await async_get_client( - username, password, session=session, use_legacy_auth=True - ) + client = await async_get_client_with_credentials(hass, username, password) except InvalidCredentialsError: errors["base"] = "invalid_auth" except NotionError as err: @@ -50,7 +56,12 @@ async def async_validate_credentials( LOGGER.exception("Unknown error while validation credentials: %s", err) errors["base"] = "unknown" - return errors + if errors: + return CredentialsValidationResult(errors=errors) + + return CredentialsValidationResult( + user_uuid=client.user_uuid, refresh_token=client.refresh_token + ) class NotionFlowHandler(config_entries.ConfigFlow, domain=DOMAIN): @@ -84,20 +95,24 @@ class NotionFlowHandler(config_entries.ConfigFlow, domain=DOMAIN): }, ) - if errors := await async_validate_credentials( + credentials_validation_result = await async_validate_credentials( self.hass, self._reauth_entry.data[CONF_USERNAME], user_input[CONF_PASSWORD] - ): + ) + + if credentials_validation_result.errors: return self.async_show_form( step_id="reauth_confirm", data_schema=REAUTH_SCHEMA, - errors=errors, + errors=credentials_validation_result.errors, description_placeholders={ CONF_USERNAME: self._reauth_entry.data[CONF_USERNAME] }, ) self.hass.config_entries.async_update_entry( - self._reauth_entry, data=self._reauth_entry.data | user_input + self._reauth_entry, + data=self._reauth_entry.data + | {CONF_REFRESH_TOKEN: credentials_validation_result.refresh_token}, ) self.hass.async_create_task( self.hass.config_entries.async_reload(self._reauth_entry.entry_id) @@ -114,13 +129,22 @@ class NotionFlowHandler(config_entries.ConfigFlow, domain=DOMAIN): await self.async_set_unique_id(user_input[CONF_USERNAME]) self._abort_if_unique_id_configured() - if errors := await async_validate_credentials( + credentials_validation_result = await async_validate_credentials( self.hass, user_input[CONF_USERNAME], user_input[CONF_PASSWORD] - ): + ) + + if credentials_validation_result.errors: return self.async_show_form( step_id="user", data_schema=AUTH_SCHEMA, - errors=errors, + errors=credentials_validation_result.errors, ) - return self.async_create_entry(title=user_input[CONF_USERNAME], data=user_input) + return self.async_create_entry( + title=user_input[CONF_USERNAME], + data={ + CONF_USERNAME: user_input[CONF_USERNAME], + CONF_USER_UUID: credentials_validation_result.user_uuid, + CONF_REFRESH_TOKEN: credentials_validation_result.refresh_token, + }, + ) diff --git a/homeassistant/components/notion/const.py b/homeassistant/components/notion/const.py index 0961b7c10c5..b1ea921a71b 100644 --- a/homeassistant/components/notion/const.py +++ b/homeassistant/components/notion/const.py @@ -4,6 +4,9 @@ import logging DOMAIN = "notion" LOGGER = logging.getLogger(__package__) +CONF_REFRESH_TOKEN = "refresh_token" +CONF_USER_UUID = "user_uuid" + SENSOR_BATTERY = "low_battery" SENSOR_DOOR = "door" SENSOR_GARAGE_DOOR = "garage_door" diff --git a/homeassistant/components/notion/diagnostics.py b/homeassistant/components/notion/diagnostics.py index 86b84760016..5c32f235639 100644 --- a/homeassistant/components/notion/diagnostics.py +++ b/homeassistant/components/notion/diagnostics.py @@ -5,12 +5,12 @@ from typing import Any from homeassistant.components.diagnostics import async_redact_data from homeassistant.config_entries import ConfigEntry -from homeassistant.const import CONF_EMAIL, CONF_PASSWORD, CONF_UNIQUE_ID, CONF_USERNAME +from homeassistant.const import CONF_EMAIL, CONF_UNIQUE_ID, CONF_USERNAME from homeassistant.core import HomeAssistant from homeassistant.helpers.update_coordinator import DataUpdateCoordinator from . import NotionData -from .const import DOMAIN +from .const import CONF_REFRESH_TOKEN, CONF_USER_UUID, DOMAIN CONF_DEVICE_KEY = "device_key" CONF_HARDWARE_ID = "hardware_id" @@ -23,12 +23,13 @@ TO_REDACT = { CONF_EMAIL, CONF_HARDWARE_ID, CONF_LAST_BRIDGE_HARDWARE_ID, - CONF_PASSWORD, + CONF_REFRESH_TOKEN, # Config entry title and unique ID may contain sensitive data: CONF_TITLE, CONF_UNIQUE_ID, CONF_USERNAME, CONF_USER_ID, + CONF_USER_UUID, } diff --git a/homeassistant/components/notion/manifest.json b/homeassistant/components/notion/manifest.json index 662114742bd..9f725587e60 100644 --- a/homeassistant/components/notion/manifest.json +++ b/homeassistant/components/notion/manifest.json @@ -7,5 +7,5 @@ "integration_type": "hub", "iot_class": "cloud_polling", "loggers": ["aionotion"], - "requirements": ["aionotion==2024.02.0"] + "requirements": ["aionotion==2024.02.1"] } diff --git a/homeassistant/components/notion/util.py b/homeassistant/components/notion/util.py new file mode 100644 index 00000000000..553199b7c7a --- /dev/null +++ b/homeassistant/components/notion/util.py @@ -0,0 +1,30 @@ +"""Define notion utilities.""" +from aionotion import ( + async_get_client_with_credentials as cwc, + async_get_client_with_refresh_token as cwrt, +) +from aionotion.client import Client + +from homeassistant.core import HomeAssistant +from homeassistant.helpers import aiohttp_client +from homeassistant.helpers.instance_id import async_get + + +async def async_get_client_with_credentials( + hass: HomeAssistant, email: str, password: str +) -> Client: + """Get a Notion client with credentials.""" + session = aiohttp_client.async_get_clientsession(hass) + instance_id = await async_get(hass) + return await cwc(email, password, session=session, session_name=instance_id) + + +async def async_get_client_with_refresh_token( + hass: HomeAssistant, user_uuid: str, refresh_token: str +) -> Client: + """Get a Notion client with credentials.""" + session = aiohttp_client.async_get_clientsession(hass) + instance_id = await async_get(hass) + return await cwrt( + user_uuid, refresh_token, session=session, session_name=instance_id + ) diff --git a/requirements_all.txt b/requirements_all.txt index af24498b092..26462fa2f3e 100644 --- a/requirements_all.txt +++ b/requirements_all.txt @@ -315,7 +315,7 @@ aiomusiccast==0.14.8 aionanoleaf==0.2.1 # homeassistant.components.notion -aionotion==2024.02.0 +aionotion==2024.02.1 # homeassistant.components.oncue aiooncue==0.3.5 diff --git a/requirements_test_all.txt b/requirements_test_all.txt index 046383af670..50abd5dba79 100644 --- a/requirements_test_all.txt +++ b/requirements_test_all.txt @@ -288,7 +288,7 @@ aiomusiccast==0.14.8 aionanoleaf==0.2.1 # homeassistant.components.notion -aionotion==2024.02.0 +aionotion==2024.02.1 # homeassistant.components.oncue aiooncue==0.3.5 diff --git a/tests/components/notion/conftest.py b/tests/components/notion/conftest.py index 3623782429f..61a54ca58ba 100644 --- a/tests/components/notion/conftest.py +++ b/tests/components/notion/conftest.py @@ -17,6 +17,8 @@ from tests.common import MockConfigEntry, load_fixture TEST_USERNAME = "user@host.com" TEST_PASSWORD = "password123" +TEST_REFRESH_TOKEN = "abcde12345" +TEST_USER_UUID = "xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx" @pytest.fixture @@ -47,6 +49,7 @@ def client_fixture(data_bridge, data_listener, data_sensor, data_user_preference ] ) ), + refresh_token=TEST_REFRESH_TOKEN, sensor=Mock( async_all=AsyncMock( return_value=[ @@ -61,6 +64,7 @@ def client_fixture(data_bridge, data_listener, data_sensor, data_user_preference ) ) ), + user_uuid=TEST_USER_UUID, ) @@ -107,7 +111,7 @@ def data_user_preferences_fixture(): @pytest.fixture(name="get_client") def get_client_fixture(client): - """Define a fixture to mock the async_get_client method.""" + """Define a fixture to mock the client retrieval methods.""" return AsyncMock(return_value=client) @@ -115,10 +119,13 @@ def get_client_fixture(client): async def mock_aionotion_fixture(client): """Define a fixture to patch aionotion.""" with patch( - "homeassistant.components.notion.async_get_client", + "homeassistant.components.notion.async_get_client_with_credentials", AsyncMock(return_value=client), ), patch( - "homeassistant.components.notion.config_flow.async_get_client", + "homeassistant.components.notion.async_get_client_with_refresh_token", + AsyncMock(return_value=client), + ), patch( + "homeassistant.components.notion.config_flow.async_get_client_with_credentials", AsyncMock(return_value=client), ): yield diff --git a/tests/components/notion/test_config_flow.py b/tests/components/notion/test_config_flow.py index e9f340fae17..646bd7a6e87 100644 --- a/tests/components/notion/test_config_flow.py +++ b/tests/components/notion/test_config_flow.py @@ -5,12 +5,12 @@ from aionotion.errors import InvalidCredentialsError, NotionError import pytest from homeassistant import data_entry_flow -from homeassistant.components.notion import DOMAIN +from homeassistant.components.notion import CONF_REFRESH_TOKEN, CONF_USER_UUID, DOMAIN from homeassistant.config_entries import SOURCE_REAUTH, SOURCE_USER from homeassistant.const import CONF_PASSWORD, CONF_USERNAME from homeassistant.core import HomeAssistant -from .conftest import TEST_PASSWORD, TEST_USERNAME +from .conftest import TEST_REFRESH_TOKEN, TEST_USER_UUID, TEST_USERNAME pytestmark = pytest.mark.usefixtures("mock_setup_entry") @@ -40,7 +40,7 @@ async def test_create_entry( # Test errors that can arise when getting a Notion API client: with patch( - "homeassistant.components.notion.config_flow.async_get_client", + "homeassistant.components.notion.config_flow.async_get_client_with_credentials", get_client_with_exception, ): result = await hass.config_entries.flow.async_init( @@ -55,8 +55,9 @@ async def test_create_entry( assert result["type"] == data_entry_flow.FlowResultType.CREATE_ENTRY assert result["title"] == TEST_USERNAME assert result["data"] == { + CONF_REFRESH_TOKEN: TEST_REFRESH_TOKEN, CONF_USERNAME: TEST_USERNAME, - CONF_PASSWORD: TEST_PASSWORD, + CONF_USER_UUID: TEST_USER_UUID, } @@ -99,7 +100,7 @@ async def test_reauth( # Test errors that can arise when getting a Notion API client: with patch( - "homeassistant.components.notion.config_flow.async_get_client", + "homeassistant.components.notion.config_flow.async_get_client_with_credentials", get_client_with_exception, ): result = await hass.config_entries.flow.async_configure( diff --git a/tests/components/notion/test_diagnostics.py b/tests/components/notion/test_diagnostics.py index a2b829281f8..f0ca64807e1 100644 --- a/tests/components/notion/test_diagnostics.py +++ b/tests/components/notion/test_diagnostics.py @@ -21,7 +21,11 @@ async def test_entry_diagnostics( "minor_version": 1, "domain": DOMAIN, "title": REDACTED, - "data": {"username": REDACTED, "password": REDACTED}, + "data": { + "refresh_token": REDACTED, + "user_uuid": REDACTED, + "username": REDACTED, + }, "options": {}, "pref_disable_new_entities": False, "pref_disable_polling": False,