Use DataUpdateCoordinator for mikrotik (#72954)
This commit is contained in:
Rami Mosleh
GitHub
parent
329ecc74c4
commit
8905e6f726
6 changed files with 81 additions and 194 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
_LOGGER = logging.getLogger(__name__)
|
||||
from .hub import MikrotikDataUpdateCoordinator
|
||||
|
||||
# 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()
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
self.available = False
|
||||
return
|
||||
except (CannotConnect, socket.timeout, OSError) as err:
|
||||
raise UpdateFailed from err
|
||||
|
||||
if not device_list:
|
||||
return
|
||||
|
|
@ -263,6 +260,7 @@ class MikrotikData:
|
|||
socket.timeout,
|
||||
) as api_error:
|
||||
_LOGGER.error("Mikrotik %s connection error %s", self._host, api_error)
|
||||
if not self.connect_to_hub():
|
||||
raise CannotConnect from api_error
|
||||
except librouteros.exceptions.ProtocolError as api_error:
|
||||
_LOGGER.warning(
|
||||
|
|
@ -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 MikrotikHub:
|
||||
class MikrotikDataUpdateCoordinator(DataUpdateCoordinator):
|
||||
"""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)
|
||||
async_dispatcher_send(self.hass, self.signal_update)
|
||||
await self.hass.async_add_executor_job(self._mk_data.update_devices)
|
||||
|
||||
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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
"homeassistant.components.mikrotik.dr.async_get",
|
||||
return_value=mock_registry,
|
||||
):
|
||||
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",
|
||||
}
|
||||
await hass.config_entries.async_setup(entry.entry_id)
|
||||
assert entry.state == ConfigEntryState.LOADED
|
||||
assert hass.data[DOMAIN][entry.entry_id]
|
||||
|
||||
|
||||
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_hub.return_value.async_setup = AsyncMock(return_value=False)
|
||||
assert await mikrotik.async_setup_entry(hass, entry) is False
|
||||
mock_api.side_effect = ConnectionClosed
|
||||
|
||||
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(
|
||||
"homeassistant.helpers.device_registry.async_get",
|
||||
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
|
||||
await hass.config_entries.async_setup(entry.entry_id)
|
||||
await hass.async_block_till_done()
|
||||
|
||||
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.entry_id not in hass.data[mikrotik.DOMAIN]
|
||||
assert entry.state == ConfigEntryState.NOT_LOADED
|
||||
assert entry.entry_id not in hass.data[DOMAIN]
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue