Store access token in entry for Fyta

homeassistant/components/fyta/__init__.py

@@ -2,15 +2,23 @@
 
 from __future__ import annotations
 
+from datetime import datetime
 import logging
+from typing import Any
+from zoneinfo import ZoneInfo
 
 from fyta_cli.fyta_connector import FytaConnector
 
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import CONF_PASSWORD, CONF_USERNAME, Platform
+from homeassistant.const import (
+    CONF_ACCESS_TOKEN,
+    CONF_PASSWORD,
+    CONF_USERNAME,
+    Platform,
+)
 from homeassistant.core import HomeAssistant
 
-from .const import DOMAIN
+from .const import CONF_EXPIRATION, DOMAIN
 from .coordinator import FytaCoordinator
 
 _LOGGER = logging.getLogger(__name__)
@@ -22,11 +30,16 @@
 
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up the Fyta integration."""
+    tz: str = hass.config.time_zone
 
     username = entry.data[CONF_USERNAME]
     password = entry.data[CONF_PASSWORD]
+    access_token: str = entry.data[CONF_ACCESS_TOKEN]
+    expiration: datetime = datetime.fromisoformat(
+        entry.data[CONF_EXPIRATION]
+    ).astimezone(ZoneInfo(tz))
 
-    fyta = FytaConnector(username, password)
+    fyta = FytaConnector(username, password, access_token, expiration, tz)
 
     coordinator = FytaCoordinator(hass, fyta)
 
@@ -47,3 +60,36 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         hass.data[DOMAIN].pop(entry.entry_id)
 
     return unload_ok
+
+
+async def async_migrate_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:
+    """Migrate old entry."""
+    _LOGGER.debug("Migrating from version %s", config_entry.version)
+
+    if config_entry.version > 1:
+        # This means the user has downgraded from a future version
+        return False
+
+    if config_entry.version == 1:
+        new = {**config_entry.data}
+        if config_entry.minor_version < 2:
+            fyta = FytaConnector(
+                config_entry.data[CONF_USERNAME], config_entry.data[CONF_PASSWORD]
+            )
+            credentials: dict[str, Any] = await fyta.login()
+            await fyta.client.close()
+
+            new[CONF_ACCESS_TOKEN] = credentials[CONF_ACCESS_TOKEN]
+            new[CONF_EXPIRATION] = credentials[CONF_EXPIRATION].isoformat()
+
+        hass.config_entries.async_update_entry(
+            config_entry, data=new, minor_version=2, version=1
+        )
+
+    _LOGGER.debug(
+        "Migration to version %s.%s successful",
+        config_entry.version,
+        config_entry.minor_version,
+    )
+
+    return True

homeassistant/components/fyta/config_flow.py

@@ -17,7 +17,7 @@
 from homeassistant.config_entries import ConfigEntry, ConfigFlow, ConfigFlowResult
 from homeassistant.const import CONF_PASSWORD, CONF_USERNAME
 
-from .const import DOMAIN
+from .const import CONF_EXPIRATION, DOMAIN
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -31,14 +31,19 @@ class FytaConfigFlow(ConfigFlow, domain=DOMAIN):
     """Handle a config flow for Fyta."""
 
     VERSION = 1
-    _entry: ConfigEntry | None = None
+    MINOR_VERSION = 2
+
+    def __init__(self) -> None:
+        """Initialize FytaConfigFlow."""
+        self.credentials: dict[str, Any] = {}
+        self._entry: ConfigEntry | None = None
 
     async def async_auth(self, user_input: Mapping[str, Any]) -> dict[str, str]:
         """Reusable Auth Helper."""
         fyta = FytaConnector(user_input[CONF_USERNAME], user_input[CONF_PASSWORD])
 
         try:
-            await fyta.login()
+            self.credentials = await fyta.login()
         except FytaConnectionError:
             return {"base": "cannot_connect"}
         except FytaAuthentificationError:
@@ -51,6 +56,10 @@ async def async_auth(self, user_input: Mapping[str, Any]) -> dict[str, str]:
         finally:
             await fyta.client.close()
 
+        self.credentials[CONF_EXPIRATION] = self.credentials[
+            CONF_EXPIRATION
+        ].isoformat()
+
         return {}
 
     async def async_step_user(
@@ -62,6 +71,7 @@ async def async_step_user(
             self._async_abort_entries_match({CONF_USERNAME: user_input[CONF_USERNAME]})
 
             if not (errors := await self.async_auth(user_input)):
+                user_input |= self.credentials
                 return self.async_create_entry(
                     title=user_input[CONF_USERNAME], data=user_input
                 )
@@ -85,6 +95,7 @@ async def async_step_reauth_confirm(
         assert self._entry is not None
 
         if user_input and not (errors := await self.async_auth(user_input)):
+            user_input |= self.credentials
             return self.async_update_reload_and_abort(
                 self._entry, data={**self._entry.data, **user_input}
             )

homeassistant/components/fyta/const.py

@@ -1,3 +1,4 @@
 """Const for fyta integration."""
 
 DOMAIN = "fyta"
+CONF_EXPIRATION = "expiration"

homeassistant/components/fyta/coordinator.py

@@ -12,10 +12,13 @@
 )
 
 from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import CONF_ACCESS_TOKEN
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
+from .const import CONF_EXPIRATION
+
 _LOGGER = logging.getLogger(__name__)
 
 
@@ -39,17 +42,33 @@ async def _async_update_data(
     ) -> dict[int, dict[str, Any]]:
         """Fetch data from API endpoint."""
 
-        if self.fyta.expiration is None or self.fyta.expiration < datetime.now():
+        if (
+            self.fyta.expiration is None
+            or self.fyta.expiration.timestamp() < datetime.now().timestamp()
+        ):
             await self.renew_authentication()
 
         return await self.fyta.update_all_plants()
 
-    async def renew_authentication(self) -> None:
+    async def renew_authentication(self) -> bool:
         """Renew access token for FYTA API."""
+        credentials: dict[str, Any] = {}
 
         try:
-            await self.fyta.login()
+            credentials = await self.fyta.login()
         except FytaConnectionError as ex:
             raise ConfigEntryNotReady from ex
         except (FytaAuthentificationError, FytaPasswordError) as ex:
             raise ConfigEntryAuthFailed from ex
+
+        new_config_entry = {**self.config_entry.data}
+        new_config_entry[CONF_ACCESS_TOKEN] = credentials[CONF_ACCESS_TOKEN]
+        new_config_entry[CONF_EXPIRATION] = credentials[CONF_EXPIRATION].isoformat()
+
+        self.hass.config_entries.async_update_entry(
+            self.config_entry, data=new_config_entry
+        )
+
+        _LOGGER.debug("Credentials successfully updated")
+
+        return True

tests/components/fyta/conftest.py

@@ -5,6 +5,11 @@
 
 import pytest
 
+from homeassistant.components.fyta.const import CONF_EXPIRATION
+from homeassistant.const import CONF_ACCESS_TOKEN
+
+from .test_config_flow import ACCESS_TOKEN, EXPIRATION
+
 
 @pytest.fixture
 def mock_fyta():
@@ -15,7 +20,26 @@ def mock_fyta():
         "homeassistant.components.fyta.config_flow.FytaConnector",
         return_value=mock_fyta_api,
     ) as mock_fyta_api:
-        mock_fyta_api.return_value.login.return_value = {}
+        mock_fyta_api.return_value.login.return_value = {
+            CONF_ACCESS_TOKEN: ACCESS_TOKEN,
+            CONF_EXPIRATION: EXPIRATION,
+        }
+        yield mock_fyta_api
+
+
+@pytest.fixture
+def mock_fyta_init():
+    """Build a fixture for the Fyta API that connects successfully and returns one device."""
+
+    mock_fyta_api = AsyncMock()
+    with patch(
+        "homeassistant.components.fyta.FytaConnector",
+        return_value=mock_fyta_api,
+    ) as mock_fyta_api:
+        mock_fyta_api.return_value.login.return_value = {
+            CONF_ACCESS_TOKEN: ACCESS_TOKEN,
+            CONF_EXPIRATION: EXPIRATION,
+        }
         yield mock_fyta_api
 
 

tests/components/fyta/test_config_flow.py

@@ -1,5 +1,6 @@
 """Test the fyta config flow."""
 
+from datetime import UTC, datetime
 from unittest.mock import AsyncMock
 
 from fyta_cli.fyta_exceptions import (
@@ -10,19 +11,21 @@
 import pytest
 
 from homeassistant import config_entries
-from homeassistant.components.fyta.const import DOMAIN
-from homeassistant.const import CONF_PASSWORD, CONF_USERNAME
+from homeassistant.components.fyta.const import CONF_EXPIRATION, DOMAIN
+from homeassistant.const import CONF_ACCESS_TOKEN, CONF_PASSWORD, CONF_USERNAME
 from homeassistant.core import HomeAssistant
 from homeassistant.data_entry_flow import FlowResultType
 
 from tests.common import MockConfigEntry
 
 USERNAME = "fyta_user"
 PASSWORD = "fyta_pass"
+ACCESS_TOKEN = "123xyz"
+EXPIRATION = datetime.fromisoformat("2024-12-31T10:00:00").astimezone(UTC)
 
 
 async def test_user_flow(
-    hass: HomeAssistant, mock_fyta: AsyncMock, mock_setup_entry
+    hass: HomeAssistant, mock_fyta: AsyncMock, mock_setup_entry: AsyncMock
 ) -> None:
     """Test we get the form."""
 
@@ -39,7 +42,12 @@ async def test_user_flow(
 
     assert result2["type"] is FlowResultType.CREATE_ENTRY
     assert result2["title"] == USERNAME
-    assert result2["data"] == {CONF_USERNAME: USERNAME, CONF_PASSWORD: PASSWORD}
+    assert result2["data"] == {
+        CONF_USERNAME: USERNAME,
+        CONF_PASSWORD: PASSWORD,
+        CONF_ACCESS_TOKEN: ACCESS_TOKEN,
+        CONF_EXPIRATION: "2024-12-31T10:00:00+00:00",
+    }
     assert len(mock_setup_entry.mock_calls) == 1
 
 
@@ -57,7 +65,7 @@ async def test_form_exceptions(
     exception: Exception,
     error: dict[str, str],
     mock_fyta: AsyncMock,
-    mock_setup_entry,
+    mock_setup_entry: AsyncMock,
 ) -> None:
     """Test we can handle Form exceptions."""
 
@@ -89,6 +97,8 @@ async def test_form_exceptions(
     assert result["title"] == USERNAME
     assert result["data"][CONF_USERNAME] == USERNAME
     assert result["data"][CONF_PASSWORD] == PASSWORD
+    assert result["data"][CONF_ACCESS_TOKEN] == ACCESS_TOKEN
+    assert result["data"][CONF_EXPIRATION] == "2024-12-31T10:00:00+00:00"
 
     assert len(mock_setup_entry.mock_calls) == 1
 
@@ -134,14 +144,19 @@ async def test_reauth(
     exception: Exception,
     error: dict[str, str],
     mock_fyta: AsyncMock,
-    mock_setup_entry,
+    mock_setup_entry: AsyncMock,
 ) -> None:
     """Test reauth-flow works."""
 
     entry = MockConfigEntry(
         domain=DOMAIN,
         title=USERNAME,
-        data={CONF_USERNAME: USERNAME, CONF_PASSWORD: PASSWORD},
+        data={
+            CONF_USERNAME: USERNAME,
+            CONF_PASSWORD: PASSWORD,
+            CONF_ACCESS_TOKEN: ACCESS_TOKEN,
+            CONF_EXPIRATION: "2024-06-30T10:00:00+00:00",
+        },
     )
     entry.add_to_hass(hass)
 
@@ -157,7 +172,8 @@ async def test_reauth(
 
     # tests with connection error
     result = await hass.config_entries.flow.async_configure(
-        result["flow_id"], {CONF_USERNAME: USERNAME, CONF_PASSWORD: PASSWORD}
+        result["flow_id"],
+        {CONF_USERNAME: USERNAME, CONF_PASSWORD: PASSWORD},
     )
     await hass.async_block_till_done()
 
@@ -178,5 +194,5 @@ async def test_reauth(
     assert result["reason"] == "reauth_successful"
     assert entry.data[CONF_USERNAME] == "other_username"
     assert entry.data[CONF_PASSWORD] == "other_password"
-
-    assert len(mock_setup_entry.mock_calls) == 1
+    assert entry.data[CONF_ACCESS_TOKEN] == ACCESS_TOKEN
+    assert entry.data[CONF_EXPIRATION] == "2024-12-31T10:00:00+00:00"

tests/components/fyta/test_init.py

@@ -0,0 +1,42 @@
+"""Test the initialization."""
+
+from unittest.mock import AsyncMock
+
+from homeassistant.components.fyta.const import CONF_EXPIRATION, DOMAIN
+from homeassistant.const import CONF_ACCESS_TOKEN, CONF_PASSWORD, CONF_USERNAME
+from homeassistant.core import HomeAssistant
+
+from .test_config_flow import ACCESS_TOKEN, PASSWORD, USERNAME
+
+from tests.common import MockConfigEntry
+
+
+async def test_migrate_config_entry(
+    hass: HomeAssistant,
+    mock_fyta_init: AsyncMock,
+) -> None:
+    """Test successful migration of entry data."""
+    entry = MockConfigEntry(
+        domain=DOMAIN,
+        title=USERNAME,
+        data={
+            CONF_USERNAME: USERNAME,
+            CONF_PASSWORD: PASSWORD,
+        },
+        version=1,
+        minor_version=1,
+    )
+    entry.add_to_hass(hass)
+
+    assert entry.version == 1
+    assert entry.minor_version == 1
+
+    await hass.config_entries.async_setup(entry.entry_id)
+    await hass.async_block_till_done()
+
+    assert entry.version == 1
+    assert entry.minor_version == 2
+    assert entry.data[CONF_USERNAME] == USERNAME
+    assert entry.data[CONF_PASSWORD] == PASSWORD
+    assert entry.data[CONF_ACCESS_TOKEN] == ACCESS_TOKEN
+    assert entry.data[CONF_EXPIRATION] == "2024-12-31T10:00:00+00:00"