Use DataUpdateCoordinator for mikrotik (#72954)

homeassistant/components/mikrotik/__init__.py

@@ -4,7 +4,7 @@ from homeassistant.core import HomeAssistant
 from homeassistant.helpers import config_validation as cv, device_registry as dr
 
 from .const import ATTR_MANUFACTURER, DOMAIN, PLATFORMS
-from .hub import MikrotikHub
+from .hub import MikrotikDataUpdateCoordinator
 
 CONFIG_SCHEMA = cv.removed(DOMAIN, raise_if_present=False)
 
@@ -12,11 +12,16 @@ CONFIG_SCHEMA = cv.removed(DOMAIN, raise_if_present=False)
 async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:
     """Set up the Mikrotik component."""
 
-    hub = MikrotikHub(hass, config_entry)
+    hub = MikrotikDataUpdateCoordinator(hass, config_entry)
     if not await hub.async_setup():
         return False
 
+    await hub.async_config_entry_first_refresh()
+
     hass.data.setdefault(DOMAIN, {})[config_entry.entry_id] = hub
+
+    hass.config_entries.async_setup_platforms(config_entry, PLATFORMS)
+
     device_registry = dr.async_get(hass)
     device_registry.async_get_or_create(
         config_entry_id=config_entry.entry_id,

homeassistant/components/mikrotik/device_tracker.py

@@ -1,8 +1,6 @@
 """Support for Mikrotik routers as device tracker."""
 from __future__ import annotations
 
-import logging
-
 from homeassistant.components.device_tracker.config_entry import ScannerEntity
 from homeassistant.components.device_tracker.const import (
     DOMAIN as DEVICE_TRACKER,
@@ -11,13 +9,12 @@ from homeassistant.components.device_tracker.const import (
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers import entity_registry
-from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
+from homeassistant.helpers.update_coordinator import CoordinatorEntity
 import homeassistant.util.dt as dt_util
 
 from .const import DOMAIN
-
+from .hub import MikrotikDataUpdateCoordinator
-_LOGGER = logging.getLogger(__name__)
 
 # These are normalized to ATTR_IP and ATTR_MAC to conform
 # to device_tracker
@@ -32,7 +29,7 @@ async def async_setup_entry(
     """Set up device tracker for Mikrotik component."""
     hub = hass.data[DOMAIN][config_entry.entry_id]
 
-    tracked: dict[str, MikrotikHubTracker] = {}
+    tracked: dict[str, MikrotikDataUpdateCoordinatorTracker] = {}
 
     registry = entity_registry.async_get(hass)
 
@@ -56,7 +53,7 @@ async def async_setup_entry(
         """Update the status of the device."""
         update_items(hub, async_add_entities, tracked)
 
-    async_dispatcher_connect(hass, hub.signal_update, update_hub)
+    config_entry.async_on_unload(hub.async_add_listener(update_hub))
 
     update_hub()
 
@@ -67,21 +64,22 @@ def update_items(hub, async_add_entities, tracked):
     new_tracked = []
     for mac, device in hub.api.devices.items():
         if mac not in tracked:
-            tracked[mac] = MikrotikHubTracker(device, hub)
+            tracked[mac] = MikrotikDataUpdateCoordinatorTracker(device, hub)
             new_tracked.append(tracked[mac])
 
     if new_tracked:
         async_add_entities(new_tracked)
 
 
-class MikrotikHubTracker(ScannerEntity):
+class MikrotikDataUpdateCoordinatorTracker(CoordinatorEntity, ScannerEntity):
     """Representation of network device."""
 
+    coordinator: MikrotikDataUpdateCoordinator
+
     def __init__(self, device, hub):
         """Initialize the tracked device."""
+        super().__init__(hub)
         self.device = device
-        self.hub = hub
-        self.unsub_dispatcher = None
 
     @property
     def is_connected(self):
@@ -89,7 +87,7 @@ class MikrotikHubTracker(ScannerEntity):
         if (
             self.device.last_seen
             and (dt_util.utcnow() - self.device.last_seen)
-            < self.hub.option_detection_time
+            < self.coordinator.option_detection_time
         ):
             return True
         return False
@@ -125,33 +123,9 @@ class MikrotikHubTracker(ScannerEntity):
         """Return a unique identifier for this device."""
         return self.device.mac
 
-    @property
-    def available(self) -> bool:
-        """Return if controller is available."""
-        return self.hub.available
-
     @property
     def extra_state_attributes(self):
         """Return the device state attributes."""
         if self.is_connected:
             return {k: v for k, v in self.device.attrs.items() if k not in FILTER_ATTRS}
         return None
-
-    async def async_added_to_hass(self):
-        """Client entity created."""
-        _LOGGER.debug("New network device tracker %s (%s)", self.name, self.unique_id)
-        self.unsub_dispatcher = async_dispatcher_connect(
-            self.hass, self.hub.signal_update, self.async_write_ha_state
-        )
-
-    async def async_update(self):
-        """Synchronize state with hub."""
-        _LOGGER.debug(
-            "Updating Mikrotik tracked client %s (%s)", self.entity_id, self.unique_id
-        )
-        await self.hub.request_update()
-
-    async def will_remove_from_hass(self):
-        """Disconnect from dispatcher."""
-        if self.unsub_dispatcher:
-            self.unsub_dispatcher()

homeassistant/components/mikrotik/hub.py

@@ -9,7 +9,7 @@ from librouteros.login import plain as login_plain, token as login_token
 
 from homeassistant.const import CONF_HOST, CONF_PASSWORD, CONF_USERNAME, CONF_VERIFY_SSL
 from homeassistant.exceptions import ConfigEntryNotReady
-from homeassistant.helpers.dispatcher import async_dispatcher_send
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 from homeassistant.util import slugify
 import homeassistant.util.dt as dt_util
 
@@ -25,13 +25,13 @@ from .const import (
     CONF_FORCE_DHCP,
     DEFAULT_DETECTION_TIME,
     DHCP,
+    DOMAIN,
     IDENTITY,
     INFO,
     IS_CAPSMAN,
     IS_WIRELESS,
     MIKROTIK_SERVICES,
     NAME,
-    PLATFORMS,
     WIRELESS,
 )
 from .errors import CannotConnect, LoginError
@@ -154,10 +154,8 @@ class MikrotikData:
         """Connect to hub."""
         try:
             self.api = get_api(self.hass, self.config_entry.data)
-            self.available = True
             return True
         except (LoginError, CannotConnect):
-            self.available = False
             return False
 
     def get_list_from_interface(self, interface):
@@ -194,9 +192,8 @@ class MikrotikData:
             # get new hub firmware version if updated
             self.firmware = self.get_info(ATTR_FIRMWARE)
 
-        except (CannotConnect, socket.timeout, OSError):
+        except (CannotConnect, socket.timeout, OSError) as err:
-            self.available = False
+            raise UpdateFailed from err
-            return
 
         if not device_list:
             return
@@ -263,7 +260,8 @@ class MikrotikData:
             socket.timeout,
         ) as api_error:
             _LOGGER.error("Mikrotik %s connection error %s", self._host, api_error)
-            raise CannotConnect from api_error
+            if not self.connect_to_hub():
+                raise CannotConnect from api_error
         except librouteros.exceptions.ProtocolError as api_error:
             _LOGGER.warning(
                 "Mikrotik %s failed to retrieve data. cmd=[%s] Error: %s",
@@ -275,15 +273,8 @@ class MikrotikData:
 
         return response if response else None
 
-    def update(self):
-        """Update device_tracker from Mikrotik API."""
-        if (not self.available or not self.api) and not self.connect_to_hub():
-            return
-        _LOGGER.debug("updating network devices for host: %s", self._host)
-        self.update_devices()
 
-
+class MikrotikDataUpdateCoordinator(DataUpdateCoordinator):
-class MikrotikHub:
     """Mikrotik Hub Object."""
 
     def __init__(self, hass, config_entry):
@@ -291,7 +282,13 @@ class MikrotikHub:
         self.hass = hass
         self.config_entry = config_entry
         self._mk_data = None
-        self.progress = None
+        super().__init__(
+            self.hass,
+            _LOGGER,
+            name=f"{DOMAIN} - {self.host}",
+            update_method=self.async_update,
+            update_interval=timedelta(seconds=10),
+        )
 
     @property
     def host(self):
@@ -328,11 +325,6 @@ class MikrotikHub:
         """Config entry option defining number of seconds from last seen to away."""
         return timedelta(seconds=self.config_entry.options[CONF_DETECTION_TIME])
 
-    @property
-    def signal_update(self):
-        """Event specific per Mikrotik entry to signal updates."""
-        return f"mikrotik-update-{self.host}"
-
     @property
     def api(self):
         """Represent Mikrotik data object."""
@@ -354,21 +346,9 @@ class MikrotikHub:
                 self.config_entry, data=data, options=options
             )
 
-    async def request_update(self):
-        """Request an update."""
-        if self.progress is not None:
-            await self.progress
-            return
-
-        self.progress = self.hass.async_create_task(self.async_update())
-        await self.progress
-
-        self.progress = None
-
     async def async_update(self):
         """Update Mikrotik devices information."""
-        await self.hass.async_add_executor_job(self._mk_data.update)
+        await self.hass.async_add_executor_job(self._mk_data.update_devices)
-        async_dispatcher_send(self.hass, self.signal_update)
 
     async def async_setup(self):
         """Set up the Mikrotik hub."""
@@ -384,9 +364,7 @@ class MikrotikHub:
         self._mk_data = MikrotikData(self.hass, self.config_entry, api)
         await self.async_add_options()
         await self.hass.async_add_executor_job(self._mk_data.get_hub_details)
-        await self.hass.async_add_executor_job(self._mk_data.update)
 
-        self.hass.config_entries.async_setup_platforms(self.config_entry, PLATFORMS)
         return True
 
 

tests/components/mikrotik/test_device_tracker.py

@@ -90,7 +90,7 @@ async def test_device_trackers(hass, mock_device_registry_devices):
         # test device_2 is added after connecting to wireless network
         WIRELESS_DATA.append(DEVICE_2_WIRELESS)
 
-        await hub.async_update()
+        await hub.async_refresh()
         await hass.async_block_till_done()
 
         device_2 = hass.states.get("device_tracker.device_2")
@@ -117,7 +117,7 @@ async def test_device_trackers(hass, mock_device_registry_devices):
         hub.api.devices["00:00:00:00:00:02"]._last_seen = dt_util.utcnow() - timedelta(
             minutes=5
         )
-        await hub.async_update()
+        await hub.async_refresh()
         await hass.async_block_till_done()
 
         device_2 = hass.states.get("device_tracker.device_2")

tests/components/mikrotik/test_hub.py

@@ -1,18 +1,7 @@
 """Test Mikrotik hub."""
 from unittest.mock import patch
 
-import librouteros
-
-from homeassistant import config_entries
 from homeassistant.components import mikrotik
-from homeassistant.const import (
-    CONF_HOST,
-    CONF_NAME,
-    CONF_PASSWORD,
-    CONF_PORT,
-    CONF_USERNAME,
-    CONF_VERIFY_SSL,
-)
 
 from . import ARP_DATA, DHCP_DATA, MOCK_DATA, MOCK_OPTIONS, WIRELESS_DATA
 
@@ -55,63 +44,6 @@ async def setup_mikrotik_entry(hass, **kwargs):
         return hass.data[mikrotik.DOMAIN][config_entry.entry_id]
 
 
-async def test_hub_setup_successful(hass):
-    """Successful setup of Mikrotik hub."""
-    with patch(
-        "homeassistant.config_entries.ConfigEntries.async_forward_entry_setup",
-        return_value=True,
-    ) as forward_entry_setup:
-        hub = await setup_mikrotik_entry(hass)
-
-    assert hub.config_entry.data == {
-        CONF_NAME: "Mikrotik",
-        CONF_HOST: "0.0.0.0",
-        CONF_USERNAME: "user",
-        CONF_PASSWORD: "pass",
-        CONF_PORT: 8278,
-        CONF_VERIFY_SSL: False,
-    }
-    assert hub.config_entry.options == {
-        mikrotik.const.CONF_FORCE_DHCP: False,
-        mikrotik.const.CONF_ARP_PING: False,
-        mikrotik.const.CONF_DETECTION_TIME: 300,
-    }
-
-    assert hub.api.available is True
-    assert hub.signal_update == "mikrotik-update-0.0.0.0"
-    assert forward_entry_setup.mock_calls[0][1] == (hub.config_entry, "device_tracker")
-
-
-async def test_hub_setup_failed(hass):
-    """Failed setup of Mikrotik hub."""
-
-    config_entry = MockConfigEntry(domain=mikrotik.DOMAIN, data=MOCK_DATA)
-    config_entry.add_to_hass(hass)
-    # error when connection fails
-    with patch(
-        "librouteros.connect", side_effect=librouteros.exceptions.ConnectionClosed
-    ):
-
-        await hass.config_entries.async_setup(config_entry.entry_id)
-
-        assert config_entry.state is config_entries.ConfigEntryState.SETUP_RETRY
-
-    # error when username or password is invalid
-    config_entry = MockConfigEntry(domain=mikrotik.DOMAIN, data=MOCK_DATA)
-    config_entry.add_to_hass(hass)
-    with patch(
-        "homeassistant.config_entries.ConfigEntries.async_forward_entry_setup"
-    ) as forward_entry_setup, patch(
-        "librouteros.connect",
-        side_effect=librouteros.exceptions.TrapError("invalid user name or password"),
-    ):
-
-        result = await hass.config_entries.async_setup(config_entry.entry_id)
-
-        assert result is False
-        assert len(forward_entry_setup.mock_calls) == 0
-
-
 async def test_update_failed(hass):
     """Test failing to connect during update."""
 
@@ -120,9 +52,9 @@ async def test_update_failed(hass):
     with patch.object(
         mikrotik.hub.MikrotikData, "command", side_effect=mikrotik.errors.CannotConnect
     ):
-        await hub.async_update()
+        await hub.async_refresh()
 
-    assert hub.api.available is False
+    assert not hub.last_update_success
 
 
 async def test_hub_not_support_wireless(hass):

tests/components/mikrotik/test_init.py

@@ -1,7 +1,12 @@
 """Test Mikrotik setup process."""
-from unittest.mock import AsyncMock, Mock, patch
+from unittest.mock import patch
+
+from librouteros.exceptions import ConnectionClosed, LibRouterosError
+import pytest
 
 from homeassistant.components import mikrotik
+from homeassistant.components.mikrotik.const import DOMAIN
+from homeassistant.config_entries import ConfigEntryState
 from homeassistant.setup import async_setup_component
 
 from . import MOCK_DATA
@@ -9,6 +14,15 @@ from . import MOCK_DATA
 from tests.common import MockConfigEntry
 
 
+@pytest.fixture(autouse=True)
+def mock_api():
+    """Mock api."""
+    with patch("librouteros.create_transport"), patch(
+        "librouteros.Api.readResponse"
+    ) as mock_api:
+        yield mock_api
+
+
 async def test_setup_with_no_config(hass):
     """Test that we do not discover anything or try to set up a hub."""
     assert await async_setup_component(hass, mikrotik.DOMAIN, {}) is True
@@ -22,37 +36,13 @@ async def test_successful_config_entry(hass):
         data=MOCK_DATA,
     )
     entry.add_to_hass(hass)
-    mock_registry = Mock()
 
-    with patch.object(mikrotik, "MikrotikHub") as mock_hub, patch(
+    await hass.config_entries.async_setup(entry.entry_id)
-        "homeassistant.components.mikrotik.dr.async_get",
+    assert entry.state == ConfigEntryState.LOADED
-        return_value=mock_registry,
+    assert hass.data[DOMAIN][entry.entry_id]
-    ):
-        mock_hub.return_value.async_setup = AsyncMock(return_value=True)
-        mock_hub.return_value.serial_num = "12345678"
-        mock_hub.return_value.model = "RB750"
-        mock_hub.return_value.hostname = "mikrotik"
-        mock_hub.return_value.firmware = "3.65"
-        assert await mikrotik.async_setup_entry(hass, entry) is True
-
-    assert len(mock_hub.mock_calls) == 2
-    p_hass, p_entry = mock_hub.mock_calls[0][1]
-
-    assert p_hass is hass
-    assert p_entry is entry
-
-    assert len(mock_registry.mock_calls) == 1
-    assert mock_registry.mock_calls[0][2] == {
-        "config_entry_id": entry.entry_id,
-        "connections": {("mikrotik", "12345678")},
-        "manufacturer": mikrotik.ATTR_MANUFACTURER,
-        "model": "RB750",
-        "name": "mikrotik",
-        "sw_version": "3.65",
-    }
 
 
-async def test_hub_fail_setup(hass):
+async def test_hub_conn_error(hass, mock_api):
     """Test that a failed setup will not store the hub."""
     entry = MockConfigEntry(
         domain=mikrotik.DOMAIN,
@@ -60,14 +50,29 @@ async def test_hub_fail_setup(hass):
     )
     entry.add_to_hass(hass)
 
-    with patch.object(mikrotik, "MikrotikHub") as mock_hub:
+    mock_api.side_effect = ConnectionClosed
-        mock_hub.return_value.async_setup = AsyncMock(return_value=False)
-        assert await mikrotik.async_setup_entry(hass, entry) is False
 
-    assert mikrotik.DOMAIN not in hass.data
+    await hass.config_entries.async_setup(entry.entry_id)
+    await hass.async_block_till_done()
+    assert entry.state == ConfigEntryState.SETUP_RETRY
 
 
-async def test_unload_entry(hass):
+async def test_hub_auth_error(hass, mock_api):
+    """Test that a failed setup will not store the hub."""
+    entry = MockConfigEntry(
+        domain=mikrotik.DOMAIN,
+        data=MOCK_DATA,
+    )
+    entry.add_to_hass(hass)
+
+    mock_api.side_effect = LibRouterosError("invalid user name or password")
+
+    await hass.config_entries.async_setup(entry.entry_id)
+    await hass.async_block_till_done()
+    assert entry.state == ConfigEntryState.SETUP_ERROR
+
+
+async def test_unload_entry(hass) -> None:
     """Test being able to unload an entry."""
     entry = MockConfigEntry(
         domain=mikrotik.DOMAIN,
@@ -75,18 +80,11 @@ async def test_unload_entry(hass):
     )
     entry.add_to_hass(hass)
 
-    with patch.object(mikrotik, "MikrotikHub") as mock_hub, patch(
+    await hass.config_entries.async_setup(entry.entry_id)
-        "homeassistant.helpers.device_registry.async_get",
+    await hass.async_block_till_done()
-        return_value=Mock(),
-    ):
-        mock_hub.return_value.async_setup = AsyncMock(return_value=True)
-        mock_hub.return_value.serial_num = "12345678"
-        mock_hub.return_value.model = "RB750"
-        mock_hub.return_value.hostname = "mikrotik"
-        mock_hub.return_value.firmware = "3.65"
-        assert await mikrotik.async_setup_entry(hass, entry) is True
 
-    assert len(mock_hub.return_value.mock_calls) == 1
+    assert await hass.config_entries.async_unload(entry.entry_id)
+    await hass.async_block_till_done()
 
-    assert await mikrotik.async_unload_entry(hass, entry)
+    assert entry.state == ConfigEntryState.NOT_LOADED
-    assert entry.entry_id not in hass.data[mikrotik.DOMAIN]
+    assert entry.entry_id not in hass.data[DOMAIN]