Add firmware update to Ubiquiti airOS (#166913)

homeassistant/components/airos/__init__.py

@@ -33,14 +33,21 @@ from homeassistant.helpers import device_registry as dr, entity_registry as er
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 
 from .const import DEFAULT_SSL, DEFAULT_VERIFY_SSL, DOMAIN, SECTION_ADVANCED_SETTINGS
-from .coordinator import AirOSConfigEntry, AirOSDataUpdateCoordinator
+from .coordinator import (
+    AirOSConfigEntry,
+    AirOSDataUpdateCoordinator,
+    AirOSFirmwareUpdateCoordinator,
+    AirOSRuntimeData,
+)
 
 _PLATFORMS: list[Platform] = [
     Platform.BINARY_SENSOR,
     Platform.BUTTON,
     Platform.SENSOR,
+    Platform.UPDATE,
 ]
 
+
 _LOGGER = logging.getLogger(__name__)
 
 
@@ -86,10 +93,20 @@ async def async_setup_entry(hass: HomeAssistant, entry: AirOSConfigEntry) -> boo
 
     airos_device = airos_class(**conn_data)
 
-    coordinator = AirOSDataUpdateCoordinator(hass, entry, device_data, airos_device)
-    await coordinator.async_config_entry_first_refresh()
+    data_coordinator = AirOSDataUpdateCoordinator(
+        hass, entry, device_data, airos_device
+    )
+    await data_coordinator.async_config_entry_first_refresh()
 
-    entry.runtime_data = coordinator
+    firmware_coordinator: AirOSFirmwareUpdateCoordinator | None = None
+    if device_data["fw_major"] >= 8:
+        firmware_coordinator = AirOSFirmwareUpdateCoordinator(hass, entry, airos_device)
+        await firmware_coordinator.async_config_entry_first_refresh()
+
+    entry.runtime_data = AirOSRuntimeData(
+        status=data_coordinator,
+        firmware=firmware_coordinator,
+    )
 
     await hass.config_entries.async_forward_entry_setups(entry, _PLATFORMS)
 

homeassistant/components/airos/binary_sensor.py

@@ -87,7 +87,7 @@ async def async_setup_entry(
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up the AirOS binary sensors from a config entry."""
-    coordinator = config_entry.runtime_data
+    coordinator = config_entry.runtime_data.status
 
     entities = [
         AirOSBinarySensor(coordinator, description)

homeassistant/components/airos/button.py

@@ -31,7 +31,9 @@ async def async_setup_entry(
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up the AirOS button from a config entry."""
-    async_add_entities([AirOSRebootButton(config_entry.runtime_data, REBOOT_BUTTON)])
+    async_add_entities(
+        [AirOSRebootButton(config_entry.runtime_data.status, REBOOT_BUTTON)]
+    )
 
 
 class AirOSRebootButton(AirOSEntity, ButtonEntity):

homeassistant/components/airos/const.py

@@ -5,6 +5,7 @@ from datetime import timedelta
 DOMAIN = "airos"
 
 SCAN_INTERVAL = timedelta(minutes=1)
+UPDATE_SCAN_INTERVAL = timedelta(days=1)
 
 MANUFACTURER = "Ubiquiti"
 

homeassistant/components/airos/coordinator.py

@@ -2,7 +2,10 @@
 
 from __future__ import annotations
 
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass
 import logging
+from typing import Any, TypeVar
 
 from airos.airos6 import AirOS6, AirOS6Data
 from airos.airos8 import AirOS8, AirOS8Data
@@ -19,20 +22,61 @@ from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryAuthFailed
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 
-from .const import DOMAIN, SCAN_INTERVAL
+from .const import DOMAIN, SCAN_INTERVAL, UPDATE_SCAN_INTERVAL
 
 _LOGGER = logging.getLogger(__name__)
 
-AirOSDeviceDetect = AirOS8 | AirOS6
-AirOSDataDetect = AirOS8Data | AirOS6Data
+type AirOSDeviceDetect = AirOS8 | AirOS6
+type AirOSDataDetect = AirOS8Data | AirOS6Data
+type AirOSUpdateData = dict[str, Any]
 
-type AirOSConfigEntry = ConfigEntry[AirOSDataUpdateCoordinator]
+type AirOSConfigEntry = ConfigEntry[AirOSRuntimeData]
+
+T = TypeVar("T", bound=AirOSDataDetect | AirOSUpdateData)
+
+
+@dataclass
+class AirOSRuntimeData:
+    """Data for AirOS config entry."""
+
+    status: AirOSDataUpdateCoordinator
+    firmware: AirOSFirmwareUpdateCoordinator | None
+
+
+async def async_fetch_airos_data(
+    airos_device: AirOSDeviceDetect,
+    update_method: Callable[[], Awaitable[T]],
+) -> T:
+    """Fetch data from AirOS device."""
+    try:
+        await airos_device.login()
+        return await update_method()
+    except AirOSConnectionAuthenticationError as err:
+        _LOGGER.exception("Error authenticating with airOS device")
+        raise ConfigEntryAuthFailed(
+            translation_domain=DOMAIN, translation_key="invalid_auth"
+        ) from err
+    except (
+        AirOSConnectionSetupError,
+        AirOSDeviceConnectionError,
+        TimeoutError,
+    ) as err:
+        _LOGGER.error("Error connecting to airOS device: %s", err)
+        raise UpdateFailed(
+            translation_domain=DOMAIN,
+            translation_key="cannot_connect",
+        ) from err
+    except AirOSDataMissingError as err:
+        _LOGGER.error("Expected data not returned by airOS device: %s", err)
+        raise UpdateFailed(
+            translation_domain=DOMAIN,
+            translation_key="error_data_missing",
+        ) from err
 
 
 class AirOSDataUpdateCoordinator(DataUpdateCoordinator[AirOSDataDetect]):
-    """Class to manage fetching AirOS data from single endpoint."""
+    """Class to manage fetching AirOS status data from single endpoint."""
 
-    airos_device: AirOSDeviceDetect
     config_entry: AirOSConfigEntry
 
     def __init__(
@@ -54,28 +98,33 @@ class AirOSDataUpdateCoordinator(DataUpdateCoordinator[AirOSDataDetect]):
         )
 
     async def _async_update_data(self) -> AirOSDataDetect:
-        """Fetch data from AirOS."""
-        try:
-            await self.airos_device.login()
-            return await self.airos_device.status()
-        except AirOSConnectionAuthenticationError as err:
-            _LOGGER.exception("Error authenticating with airOS device")
-            raise ConfigEntryAuthFailed(
-                translation_domain=DOMAIN, translation_key="invalid_auth"
-            ) from err
-        except (
-            AirOSConnectionSetupError,
-            AirOSDeviceConnectionError,
-            TimeoutError,
-        ) as err:
-            _LOGGER.error("Error connecting to airOS device: %s", err)
-            raise UpdateFailed(
-                translation_domain=DOMAIN,
-                translation_key="cannot_connect",
-            ) from err
-        except AirOSDataMissingError as err:
-            _LOGGER.error("Expected data not returned by airOS device: %s", err)
-            raise UpdateFailed(
-                translation_domain=DOMAIN,
-                translation_key="error_data_missing",
-            ) from err
+        """Fetch status data from AirOS."""
+        return await async_fetch_airos_data(self.airos_device, self.airos_device.status)
+
+
+class AirOSFirmwareUpdateCoordinator(DataUpdateCoordinator[AirOSUpdateData]):
+    """Class to manage fetching AirOS firmware."""
+
+    config_entry: AirOSConfigEntry
+
+    def __init__(
+        self,
+        hass: HomeAssistant,
+        config_entry: AirOSConfigEntry,
+        airos_device: AirOSDeviceDetect,
+    ) -> None:
+        """Initialize the coordinator."""
+        self.airos_device = airos_device
+        super().__init__(
+            hass,
+            _LOGGER,
+            config_entry=config_entry,
+            name=DOMAIN,
+            update_interval=UPDATE_SCAN_INTERVAL,
+        )
+
+    async def _async_update_data(self) -> AirOSUpdateData:
+        """Fetch firmware data from AirOS."""
+        return await async_fetch_airos_data(
+            self.airos_device, self.airos_device.update_check
+        )

homeassistant/components/airos/diagnostics.py

@@ -29,5 +29,15 @@ async def async_get_config_entry_diagnostics(
     """Return diagnostics for a config entry."""
     return {
         "entry_data": async_redact_data(entry.data, TO_REDACT_HA),
-        "data": async_redact_data(entry.runtime_data.data.to_dict(), TO_REDACT_AIROS),
+        "data": {
+            "status_data": async_redact_data(
+                entry.runtime_data.status.data.to_dict(), TO_REDACT_AIROS
+            ),
+            "firmware_data": async_redact_data(
+                entry.runtime_data.firmware.data
+                if entry.runtime_data.firmware is not None
+                else {},
+                TO_REDACT_AIROS,
+            ),
+        },
     }

homeassistant/components/airos/sensor.py

@@ -180,7 +180,7 @@ async def async_setup_entry(
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up the AirOS sensors from a config entry."""
-    coordinator = config_entry.runtime_data
+    coordinator = config_entry.runtime_data.status
 
     entities = [AirOSSensor(coordinator, description) for description in COMMON_SENSORS]
 

homeassistant/components/airos/strings.json

@@ -206,6 +206,12 @@
     },
     "reboot_failed": {
       "message": "The device did not accept the reboot request. Try again, or check your device web interface for errors."
+    },
+    "update_connection_authentication_error": {
+      "message": "Authentication or connection failed during firmware update"
+    },
+    "update_error": {
+      "message": "Connection failed during firmware update"
     }
   }
 }

homeassistant/components/airos/update.py

@@ -0,0 +1,101 @@
+"""AirOS update component for Home Assistant."""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from airos.exceptions import AirOSConnectionAuthenticationError, AirOSException
+
+from homeassistant.components.update import (
+    UpdateDeviceClass,
+    UpdateEntity,
+    UpdateEntityFeature,
+)
+from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import HomeAssistantError
+from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
+
+from .const import DOMAIN
+from .coordinator import (
+    AirOSConfigEntry,
+    AirOSDataUpdateCoordinator,
+    AirOSFirmwareUpdateCoordinator,
+)
+from .entity import AirOSEntity
+
+PARALLEL_UPDATES = 0
+
+_LOGGER = logging.getLogger(__name__)
+
+
+async def async_setup_entry(
+    hass: HomeAssistant,
+    config_entry: AirOSConfigEntry,
+    async_add_entities: AddConfigEntryEntitiesCallback,
+) -> None:
+    """Set up the AirOS update entity from a config entry."""
+    runtime_data = config_entry.runtime_data
+
+    if runtime_data.firmware is None:  # Unsupported device
+        return
+    async_add_entities([AirOSUpdateEntity(runtime_data.status, runtime_data.firmware)])
+
+
+class AirOSUpdateEntity(AirOSEntity, UpdateEntity):
+    """Update entity for AirOS firmware updates."""
+
+    _attr_device_class = UpdateDeviceClass.FIRMWARE
+    _attr_supported_features = UpdateEntityFeature.INSTALL
+
+    def __init__(
+        self,
+        status: AirOSDataUpdateCoordinator,
+        firmware: AirOSFirmwareUpdateCoordinator,
+    ) -> None:
+        """Initialize the AirOS update entity."""
+        super().__init__(status)
+        self.status = status
+        self.firmware = firmware
+
+        self._attr_unique_id = f"{status.data.derived.mac}_firmware_update"
+
+    @property
+    def installed_version(self) -> str | None:
+        """Return the installed firmware version."""
+        return self.status.data.host.fwversion
+
+    @property
+    def latest_version(self) -> str | None:
+        """Return the latest firmware version."""
+        if not self.firmware.data.get("update", False):
+            return self.status.data.host.fwversion
+        return self.firmware.data.get("version")
+
+    @property
+    def release_url(self) -> str | None:
+        """Return the release url of the latest firmware."""
+        return self.firmware.data.get("changelog")
+
+    async def async_install(
+        self,
+        version: str | None,
+        backup: bool,
+        **kwargs: Any,
+    ) -> None:
+        """Handle the firmware update installation."""
+        _LOGGER.debug("Starting firmware update")
+        try:
+            await self.status.airos_device.login()
+            await self.status.airos_device.download()
+            await self.status.airos_device.install()
+        except AirOSConnectionAuthenticationError as err:
+            raise HomeAssistantError(
+                translation_domain=DOMAIN,
+                translation_key="update_connection_authentication_error",
+            ) from err
+        except AirOSException as err:
+            raise HomeAssistantError(
+                translation_domain=DOMAIN,
+                translation_key="update_error",
+            ) from err

tests/components/airos/conftest.py

@@ -1,6 +1,7 @@
 """Common fixtures for the Ubiquiti airOS tests."""
 
 from collections.abc import Generator
+from typing import Any
 from unittest.mock import AsyncMock, MagicMock, patch
 
 from airos.airos6 import AirOS6Data
@@ -9,6 +10,7 @@ from airos.helpers import DetectDeviceData
 import pytest
 
 from homeassistant.components.airos.const import DEFAULT_USERNAME, DOMAIN
+from homeassistant.components.airos.coordinator import AirOSUpdateData
 from homeassistant.const import CONF_HOST, CONF_PASSWORD, CONF_USERNAME
 
 from . import AirOSData
@@ -17,7 +19,16 @@ from tests.common import MockConfigEntry, load_json_object_fixture
 
 
 @pytest.fixture
-def ap_fixture(request: pytest.FixtureRequest) -> AirOSData:
+def ap_firmware_fixture(request: pytest.FixtureRequest) -> AirOSUpdateData:
+    """Return fixture for AP firmware data."""
+    available = getattr(request, "param", False)
+    if available:
+        return load_json_object_fixture("firmware_update_available.json", DOMAIN)
+    return load_json_object_fixture("firmware_update_latest.json", DOMAIN)
+
+
+@pytest.fixture
+def ap_status_fixture(request: pytest.FixtureRequest) -> AirOSData:
     """Load fixture data for airOS device."""
     json_data = load_json_object_fixture("airos_loco5ac_ap-ptp.json", DOMAIN)
     if hasattr(request, "param"):
@@ -61,11 +72,14 @@ def mock_airos_class() -> Generator[MagicMock]:
 
 @pytest.fixture
 def mock_airos_client(
-    mock_airos_class: MagicMock, ap_fixture: AirOSData
+    mock_airos_class: MagicMock,
+    ap_status_fixture: AirOSData,
+    ap_firmware_fixture: dict[str, Any],
 ) -> Generator[AsyncMock]:
     """Fixture to mock the AirOS API client."""
     client = mock_airos_class.return_value
-    client.status.return_value = ap_fixture
+    client.status.return_value = ap_status_fixture
+    client.update_check.return_value = ap_firmware_fixture
     client.login.return_value = True
     client.reboot.return_value = True
     return client
@@ -97,13 +111,13 @@ def mock_discovery_method() -> Generator[AsyncMock]:
 
 
 @pytest.fixture
-def mock_async_get_firmware_data(ap_fixture: AirOSData):
+def mock_async_get_firmware_data(ap_status_fixture: AirOSData):
     """Fixture to mock async_get_firmware_data to not do a network call."""
-    fw_major = int(ap_fixture.host.fwversion.lstrip("v").split(".", 1)[0])
+    fw_major = int(ap_status_fixture.host.fwversion.lstrip("v").split(".", 1)[0])
     return_value = DetectDeviceData(
         fw_major=fw_major,
-        mac=ap_fixture.derived.mac,
-        hostname=ap_fixture.host.hostname,
+        mac=ap_status_fixture.derived.mac,
+        hostname=ap_status_fixture.host.hostname,
     )
 
     mock = AsyncMock(return_value=return_value)

tests/components/airos/fixtures/firmware_update_available.json

@@ -0,0 +1,9 @@
+{
+  "checksum": "95915b6f040fedd05033f514427e99a1",
+  "version": "v8.7.22",
+  "security": "",
+  "date": "260227",
+  "url": "https://dl.ubnt.com/firmwares/XC-fw/V8.7.22/WA.v8.7.22.48486.260227.1959.bin",
+  "update": true,
+  "changelog": "https://dl.ubnt.com/firmwares/XC-fw/V8.7.22/changelog.txt"
+}

tests/components/airos/fixtures/firmware_update_latest.json

@@ -0,0 +1,3 @@
+{
+  "update": false
+}

tests/components/airos/test_binary_sensor.py

@@ -15,7 +15,7 @@ from tests.common import MockConfigEntry, snapshot_platform
 
 
 @pytest.mark.parametrize(
-    ("ap_fixture"),
+    ("ap_status_fixture"),
     [
         "airos_loco5ac_ap-ptp.json",  # v8 ptp
         "airos_liteapgps_ap_ptmp_40mhz.json",  # v8 ptmp

tests/components/airos/test_config_flow.py

@@ -84,7 +84,7 @@ MOCK_DISC_EXISTS = {
 
 async def test_manual_flow_creates_entry(
     hass: HomeAssistant,
-    ap_fixture: dict[str, Any],
+    ap_status_fixture: dict[str, Any],
     mock_airos_client: AsyncMock,
     mock_async_get_firmware_data: AsyncMock,
     mock_setup_entry: AsyncMock,
@@ -159,7 +159,7 @@ async def test_form_duplicate_entry(
 async def test_form_exception_handling(
     hass: HomeAssistant,
     mock_setup_entry: AsyncMock,
-    ap_fixture: dict[str, Any],
+    ap_status_fixture: dict[str, Any],
     mock_airos_client: AsyncMock,
     mock_async_get_firmware_data: AsyncMock,
     exception: Exception,
@@ -186,11 +186,11 @@ async def test_form_exception_handling(
     assert result["type"] is FlowResultType.FORM
     assert result["errors"] == {"base": error}
 
-    fw_major = int(ap_fixture.host.fwversion.lstrip("v").split(".", 1)[0])
+    fw_major = int(ap_status_fixture.host.fwversion.lstrip("v").split(".", 1)[0])
     valid_data = DetectDeviceData(
         fw_major=fw_major,
-        mac=ap_fixture.derived.mac,
-        hostname=ap_fixture.host.hostname,
+        mac=ap_status_fixture.derived.mac,
+        hostname=ap_status_fixture.host.hostname,
     )
 
     with patch(
@@ -210,7 +210,7 @@ async def test_form_exception_handling(
 
 async def test_reauth_flow_scenario(
     hass: HomeAssistant,
-    ap_fixture: AirOSData,
+    ap_status_fixture: AirOSData,
     mock_airos_client: AsyncMock,
     mock_config_entry: MockConfigEntry,
     mock_setup_entry: AsyncMock,
@@ -234,11 +234,11 @@ async def test_reauth_flow_scenario(
     assert flow["type"] == FlowResultType.FORM
     assert flow["step_id"] == REAUTH_STEP
 
-    fw_major = int(ap_fixture.host.fwversion.lstrip("v").split(".", 1)[0])
+    fw_major = int(ap_status_fixture.host.fwversion.lstrip("v").split(".", 1)[0])
     valid_data = DetectDeviceData(
         fw_major=fw_major,
-        mac=ap_fixture.derived.mac,
-        hostname=ap_fixture.host.hostname,
+        mac=ap_status_fixture.derived.mac,
+        hostname=ap_status_fixture.host.hostname,
     )
 
     mock_firmware = AsyncMock(return_value=valid_data)
@@ -283,7 +283,7 @@ async def test_reauth_flow_scenario(
 )
 async def test_reauth_flow_scenarios(
     hass: HomeAssistant,
-    ap_fixture: AirOSData,
+    ap_status_fixture: AirOSData,
     expected_error: str,
     mock_airos_client: AsyncMock,
     mock_async_get_firmware_data: AsyncMock,
@@ -321,11 +321,11 @@ async def test_reauth_flow_scenarios(
         assert result["step_id"] == REAUTH_STEP
         assert result["errors"] == {"base": expected_error}
 
-    fw_major = int(ap_fixture.host.fwversion.lstrip("v").split(".", 1)[0])
+    fw_major = int(ap_status_fixture.host.fwversion.lstrip("v").split(".", 1)[0])
     valid_data = DetectDeviceData(
         fw_major=fw_major,
-        mac=ap_fixture.derived.mac,
-        hostname=ap_fixture.host.hostname,
+        mac=ap_status_fixture.derived.mac,
+        hostname=ap_status_fixture.host.hostname,
     )
 
     with patch(
@@ -346,7 +346,7 @@ async def test_reauth_flow_scenarios(
 
 async def test_reauth_unique_id_mismatch(
     hass: HomeAssistant,
-    ap_fixture: AirOSData,
+    ap_status_fixture: AirOSData,
     mock_airos_client: AsyncMock,
     mock_async_get_firmware_data: AsyncMock,
     mock_config_entry: MockConfigEntry,
@@ -366,11 +366,11 @@ async def test_reauth_unique_id_mismatch(
             data=mock_config_entry.data,
         )
 
-    fw_major = int(ap_fixture.host.fwversion.lstrip("v").split(".", 1)[0])
+    fw_major = int(ap_status_fixture.host.fwversion.lstrip("v").split(".", 1)[0])
     valid_data = DetectDeviceData(
         fw_major=fw_major,
         mac="FF:23:45:67:89:AB",
-        hostname=ap_fixture.host.hostname,
+        hostname=ap_status_fixture.host.hostname,
     )
 
     with patch(
@@ -501,7 +501,7 @@ async def test_reconfigure_flow_failure(
 
 async def test_reconfigure_unique_id_mismatch(
     hass: HomeAssistant,
-    ap_fixture: AirOSData,
+    ap_status_fixture: AirOSData,
     mock_airos_client: AsyncMock,
     mock_async_get_firmware_data: AsyncMock,
     mock_config_entry: MockConfigEntry,
@@ -516,11 +516,11 @@ async def test_reconfigure_unique_id_mismatch(
     )
     flow_id = result["flow_id"]
 
-    fw_major = int(ap_fixture.host.fwversion.lstrip("v").split(".", 1)[0])
+    fw_major = int(ap_status_fixture.host.fwversion.lstrip("v").split(".", 1)[0])
     mismatched_data = DetectDeviceData(
         fw_major=fw_major,
         mac="FF:23:45:67:89:AB",
-        hostname=ap_fixture.host.hostname,
+        hostname=ap_status_fixture.host.hostname,
     )
 
     user_input = {

tests/components/airos/test_diagnostics.py

@@ -11,7 +11,7 @@ from . import setup_integration
 
 from tests.common import MockConfigEntry
 from tests.components.diagnostics import get_diagnostics_for_config_entry
-from tests.typing import ClientSessionGenerator
+from tests.typing import Any, ClientSessionGenerator
 
 
 async def test_diagnostics(
@@ -19,7 +19,8 @@ async def test_diagnostics(
     hass_client: ClientSessionGenerator,
     mock_airos_client: MagicMock,
     mock_config_entry: MockConfigEntry,
-    ap_fixture: AirOS8Data,
+    ap_status_fixture: AirOS8Data,
+    ap_firmware_fixture: dict[str, Any],
     snapshot: SnapshotAssertion,
     mock_async_get_firmware_data: AsyncMock,
 ) -> None:

tests/components/airos/test_init.py

@@ -18,9 +18,14 @@ from homeassistant.components.airos.const import (
     DOMAIN,
     SECTION_ADVANCED_SETTINGS,
 )
+from homeassistant.components.airos.coordinator import async_fetch_airos_data
 from homeassistant.components.binary_sensor import DOMAIN as BINARY_SENSOR_DOMAIN
 from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN
-from homeassistant.config_entries import SOURCE_USER, ConfigEntryState
+from homeassistant.config_entries import (
+    SOURCE_USER,
+    ConfigEntryAuthFailed,
+    ConfigEntryState,
+)
 from homeassistant.const import (
     CONF_HOST,
     CONF_PASSWORD,
@@ -282,3 +287,11 @@ async def test_setup_entry_failure(
     result = await hass.config_entries.async_setup(mock_config_entry.entry_id)
     assert result is False
     assert mock_config_entry.state == state
+
+
+async def test_fetch_airos_data_auth_error(mock_airos_client: MagicMock) -> None:
+    """Test login auth error triggers ConfigEntryAuthFailed."""
+    mock_airos_client.login.side_effect = AirOSConnectionAuthenticationError
+
+    with pytest.raises(ConfigEntryAuthFailed):
+        await async_fetch_airos_data(mock_airos_client, mock_airos_client.status)

tests/components/airos/test_update.py

@@ -0,0 +1,146 @@
+"""Test the Ubiquiti airOS firmware update."""
+
+from unittest.mock import AsyncMock
+
+from airos.exceptions import (
+    AirOSConnectionAuthenticationError,
+    AirOSDeviceConnectionError,
+    AirOSException,
+)
+import pytest
+
+from homeassistant.const import Platform
+from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import HomeAssistantError
+from homeassistant.helpers import entity_registry as er
+
+from . import setup_integration
+
+from tests.common import MockConfigEntry
+
+
+@pytest.mark.parametrize(
+    ("ap_status_fixture", "ap_firmware_fixture", "entity_id"),
+    [
+        ("airos_loco5ac_ap-ptp.json", True, "update.nanostation_5ac_ap_name_firmware"),
+        ("airos_liteapgps_ap_ptmp_40mhz.json", True, "update.house_bridge_firmware"),
+    ],
+    indirect=["ap_status_fixture", "ap_firmware_fixture"],
+)
+async def test_update_entity(
+    hass: HomeAssistant,
+    mock_config_entry: MockConfigEntry,
+    entity_registry: er.EntityRegistry,
+    mock_airos_client: AsyncMock,
+    mock_async_get_firmware_data: AsyncMock,
+    entity_id: str,
+) -> None:
+    """Test the firmware update entity behavior."""
+    await setup_integration(hass, mock_config_entry, [Platform.UPDATE])
+
+    entries = er.async_entries_for_config_entry(
+        entity_registry, mock_config_entry.entry_id
+    )
+    update_entities = [e for e in entries if e.domain == "update"]
+
+    assert len(update_entities) == 1
+    assert update_entities[0].entity_id == entity_id
+
+    state = hass.states.get(entity_id)
+    assert state is not None
+
+    await hass.services.async_call(
+        "update",
+        "install",
+        {"entity_id": entity_id},
+        blocking=True,
+    )
+
+    mock_airos_client.update_check.assert_awaited()
+    mock_airos_client.download.assert_awaited()
+    mock_airos_client.install.assert_awaited()
+
+    await hass.async_block_till_done()
+    new_state = hass.states.get(entity_id)
+    assert new_state is not None
+
+
+@pytest.mark.parametrize(
+    "ap_status_fixture",
+    [
+        "airos_NanoStation_loco_M5_v6.3.16_XM_sta.json",
+        "airos_NanoStation_M5_sta_v6.3.16.json",
+    ],
+    indirect=True,
+)
+async def test_no_update_entity(
+    hass: HomeAssistant,
+    mock_config_entry: MockConfigEntry,
+    entity_registry: er.EntityRegistry,
+    mock_airos_client: AsyncMock,
+    mock_async_get_firmware_data: AsyncMock,
+    ap_firmware_fixture: dict[str, bool],
+) -> None:
+    """Test the firmware update entity behavior is not implemented."""
+    await setup_integration(hass, mock_config_entry, [Platform.UPDATE])
+
+    entries = er.async_entries_for_config_entry(
+        entity_registry, mock_config_entry.entry_id
+    )
+    update_entities = [e for e in entries if e.domain == "update"]
+
+    assert not update_entities
+
+
+@pytest.mark.parametrize(
+    ("ap_status_fixture", "ap_firmware_fixture", "exception", "translation_key"),
+    [
+        (
+            "airos_loco5ac_ap-ptp.json",
+            True,
+            AirOSConnectionAuthenticationError,
+            "update_connection_authentication_error",
+        ),
+        (
+            "airos_liteapgps_ap_ptmp_40mhz.json",
+            True,
+            AirOSDeviceConnectionError,
+            "update_error",
+        ),
+    ],
+    indirect=["ap_status_fixture", "ap_firmware_fixture"],
+)
+async def test_update_failure(
+    hass: HomeAssistant,
+    mock_config_entry: MockConfigEntry,
+    entity_registry: er.EntityRegistry,
+    mock_airos_client: AsyncMock,
+    mock_async_get_firmware_data: AsyncMock,
+    exception: AirOSException,
+    translation_key: str,
+) -> None:
+    """Test the firmware update entity behavior."""
+    await setup_integration(hass, mock_config_entry, [Platform.UPDATE])
+
+    entries = er.async_entries_for_config_entry(
+        entity_registry, mock_config_entry.entry_id
+    )
+    update_entities = [e for e in entries if e.domain == "update"]
+
+    assert len(update_entities) == 1
+    entity_id = update_entities[0].entity_id
+
+    state = hass.states.get(entity_id)
+    assert state is not None
+
+    mock_airos_client.download.side_effect = exception
+
+    with pytest.raises(HomeAssistantError) as exc:
+        await hass.services.async_call(
+            "update",
+            "install",
+            {"entity_id": entity_id},
+            blocking=True,
+        )
+
+    assert exc.value.translation_key == translation_key