Use DataUpdateCoordinator for mikrotik (#72954)

This commit is contained in:
Rami Mosleh 2022-06-29 16:32:29 +03:00 committed by GitHub
parent 329ecc74c4
commit 8905e6f726
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
6 changed files with 81 additions and 194 deletions

View File

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

View File

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

View File

@ -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.const import CONF_HOST, CONF_PASSWORD, CONF_USERNAME, CONF_VERIFY_SSL
from homeassistant.exceptions import ConfigEntryNotReady 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 from homeassistant.util import slugify
import homeassistant.util.dt as dt_util import homeassistant.util.dt as dt_util
@ -25,13 +25,13 @@ from .const import (
CONF_FORCE_DHCP, CONF_FORCE_DHCP,
DEFAULT_DETECTION_TIME, DEFAULT_DETECTION_TIME,
DHCP, DHCP,
DOMAIN,
IDENTITY, IDENTITY,
INFO, INFO,
IS_CAPSMAN, IS_CAPSMAN,
IS_WIRELESS, IS_WIRELESS,
MIKROTIK_SERVICES, MIKROTIK_SERVICES,
NAME, NAME,
PLATFORMS,
WIRELESS, WIRELESS,
) )
from .errors import CannotConnect, LoginError from .errors import CannotConnect, LoginError
@ -154,10 +154,8 @@ class MikrotikData:
"""Connect to hub.""" """Connect to hub."""
try: try:
self.api = get_api(self.hass, self.config_entry.data) self.api = get_api(self.hass, self.config_entry.data)
self.available = True
return True return True
except (LoginError, CannotConnect): except (LoginError, CannotConnect):
self.available = False
return False return False
def get_list_from_interface(self, interface): def get_list_from_interface(self, interface):
@ -194,9 +192,8 @@ class MikrotikData:
# get new hub firmware version if updated # get new hub firmware version if updated
self.firmware = self.get_info(ATTR_FIRMWARE) 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: if not device_list:
return return
@ -263,7 +260,8 @@ class MikrotikData:
socket.timeout, socket.timeout,
) as api_error: ) as api_error:
_LOGGER.error("Mikrotik %s connection error %s", self._host, 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: except librouteros.exceptions.ProtocolError as api_error:
_LOGGER.warning( _LOGGER.warning(
"Mikrotik %s failed to retrieve data. cmd=[%s] Error: %s", "Mikrotik %s failed to retrieve data. cmd=[%s] Error: %s",
@ -275,15 +273,8 @@ class MikrotikData:
return response if response else None 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.""" """Mikrotik Hub Object."""
def __init__(self, hass, config_entry): def __init__(self, hass, config_entry):
@ -291,7 +282,13 @@ class MikrotikHub:
self.hass = hass self.hass = hass
self.config_entry = config_entry self.config_entry = config_entry
self._mk_data = None 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 @property
def host(self): def host(self):
@ -328,11 +325,6 @@ class MikrotikHub:
"""Config entry option defining number of seconds from last seen to away.""" """Config entry option defining number of seconds from last seen to away."""
return timedelta(seconds=self.config_entry.options[CONF_DETECTION_TIME]) 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 @property
def api(self): def api(self):
"""Represent Mikrotik data object.""" """Represent Mikrotik data object."""
@ -354,21 +346,9 @@ class MikrotikHub:
self.config_entry, data=data, options=options 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): async def async_update(self):
"""Update Mikrotik devices information.""" """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): async def async_setup(self):
"""Set up the Mikrotik hub.""" """Set up the Mikrotik hub."""
@ -384,9 +364,7 @@ class MikrotikHub:
self._mk_data = MikrotikData(self.hass, self.config_entry, api) self._mk_data = MikrotikData(self.hass, self.config_entry, api)
await self.async_add_options() 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.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 return True

View File

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

View File

@ -1,18 +1,7 @@
"""Test Mikrotik hub.""" """Test Mikrotik hub."""
from unittest.mock import patch from unittest.mock import patch
import librouteros
from homeassistant import config_entries
from homeassistant.components import mikrotik 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 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] 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): async def test_update_failed(hass):
"""Test failing to connect during update.""" """Test failing to connect during update."""
@ -120,9 +52,9 @@ async def test_update_failed(hass):
with patch.object( with patch.object(
mikrotik.hub.MikrotikData, "command", side_effect=mikrotik.errors.CannotConnect 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): async def test_hub_not_support_wireless(hass):

View File

@ -1,7 +1,12 @@
"""Test Mikrotik setup process.""" """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 import mikrotik
from homeassistant.components.mikrotik.const import DOMAIN
from homeassistant.config_entries import ConfigEntryState
from homeassistant.setup import async_setup_component from homeassistant.setup import async_setup_component
from . import MOCK_DATA from . import MOCK_DATA
@ -9,6 +14,15 @@ from . import MOCK_DATA
from tests.common import MockConfigEntry 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): async def test_setup_with_no_config(hass):
"""Test that we do not discover anything or try to set up a hub.""" """Test that we do not discover anything or try to set up a hub."""
assert await async_setup_component(hass, mikrotik.DOMAIN, {}) is True assert await async_setup_component(hass, mikrotik.DOMAIN, {}) is True
@ -22,37 +36,13 @@ async def test_successful_config_entry(hass):
data=MOCK_DATA, data=MOCK_DATA,
) )
entry.add_to_hass(hass) 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.""" """Test that a failed setup will not store the hub."""
entry = MockConfigEntry( entry = MockConfigEntry(
domain=mikrotik.DOMAIN, domain=mikrotik.DOMAIN,
@ -60,14 +50,29 @@ async def test_hub_fail_setup(hass):
) )
entry.add_to_hass(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.""" """Test being able to unload an entry."""
entry = MockConfigEntry( entry = MockConfigEntry(
domain=mikrotik.DOMAIN, domain=mikrotik.DOMAIN,
@ -75,18 +80,11 @@ async def test_unload_entry(hass):
) )
entry.add_to_hass(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]