From 35900964cbe8088cc7b2af297b3417c86037e46e Mon Sep 17 00:00:00 2001 From: Robert Svensson Date: Tue, 30 Jul 2019 10:05:51 +0200 Subject: [PATCH] UniFi - Track devices (#25570) --- homeassistant/components/unifi/__init__.py | 7 +- homeassistant/components/unifi/const.py | 2 + .../components/unifi/device_tracker.py | 109 ++++++++++++++++-- homeassistant/components/unifi/manifest.json | 2 +- requirements_all.txt | 2 +- requirements_test_all.txt | 2 +- tests/components/unifi/test_device_tracker.py | 55 ++++++--- tests/components/unifi/test_init.py | 2 +- 8 files changed, 147 insertions(+), 34 deletions(-) diff --git a/homeassistant/components/unifi/__init__.py b/homeassistant/components/unifi/__init__.py index 883e00a5559..627e23f4ce6 100644 --- a/homeassistant/components/unifi/__init__.py +++ b/homeassistant/components/unifi/__init__.py @@ -7,8 +7,9 @@ from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC import homeassistant.helpers.config_validation as cv from .const import ( - CONF_BLOCK_CLIENT, CONF_CONTROLLER, CONF_DETECTION_TIME, CONF_SITE_ID, - CONF_SSID_FILTER, CONTROLLER_ID, DOMAIN, UNIFI_CONFIG) + ATTR_MANUFACTURER, CONF_BLOCK_CLIENT, CONF_CONTROLLER, + CONF_DETECTION_TIME, CONF_SITE_ID, CONF_SSID_FILTER, CONTROLLER_ID, + DOMAIN, UNIFI_CONFIG) from .controller import UniFiController CONF_CONTROLLERS = 'controllers' @@ -66,7 +67,7 @@ async def async_setup_entry(hass, config_entry): device_registry.async_get_or_create( config_entry_id=config_entry.entry_id, connections={(CONNECTION_NETWORK_MAC, controller.mac)}, - manufacturer='Ubiquiti', + manufacturer=ATTR_MANUFACTURER, model="UniFi Controller", name="UniFi Controller", # sw_version=config.raw['swversion'], diff --git a/homeassistant/components/unifi/const.py b/homeassistant/components/unifi/const.py index e6076829240..017b9d76418 100644 --- a/homeassistant/components/unifi/const.py +++ b/homeassistant/components/unifi/const.py @@ -14,3 +14,5 @@ UNIFI_CONFIG = 'unifi_config' CONF_BLOCK_CLIENT = 'block_client' CONF_DETECTION_TIME = 'detection_time' CONF_SSID_FILTER = 'ssid_filter' + +ATTR_MANUFACTURER = 'Ubiquiti Networks' diff --git a/homeassistant/components/unifi/device_tracker.py b/homeassistant/components/unifi/device_tracker.py index 8c4e54e0129..3cc666c8bb5 100644 --- a/homeassistant/components/unifi/device_tracker.py +++ b/homeassistant/components/unifi/device_tracker.py @@ -20,8 +20,8 @@ import homeassistant.helpers.config_validation as cv import homeassistant.util.dt as dt_util from .const import ( - CONF_CONTROLLER, CONF_DETECTION_TIME, CONF_SITE_ID, CONF_SSID_FILTER, - CONTROLLER_ID, DOMAIN as UNIFI_DOMAIN) + ATTR_MANUFACTURER, CONF_CONTROLLER, CONF_DETECTION_TIME, CONF_SITE_ID, + CONF_SSID_FILTER, CONTROLLER_ID, DOMAIN as UNIFI_DOMAIN) LOGGER = logging.getLogger(__name__) @@ -87,7 +87,7 @@ async def async_setup_entry(hass, config_entry, async_add_entities): for entity in registry.entities.values(): if entity.config_entry_id == config_entry.entry_id and \ - entity.domain == DOMAIN: + entity.domain == DOMAIN and '-' in entity.unique_id: mac, _ = entity.unique_id.split('-', 1) @@ -116,7 +116,7 @@ def update_items(controller, async_add_entities, tracked): for client_id in controller.api.clients: if client_id in tracked: - LOGGER.debug("Updating UniFi tracked device %s (%s)", + LOGGER.debug("Updating UniFi tracked client %s (%s)", tracked[client_id].entity_id, tracked[client_id].client.mac) tracked[client_id].async_schedule_update_ha_state() @@ -131,17 +131,34 @@ def update_items(controller, async_add_entities, tracked): tracked[client_id] = UniFiClientTracker(client, controller) new_tracked.append(tracked[client_id]) - LOGGER.debug("New UniFi switch %s (%s)", client.hostname, client.mac) + LOGGER.debug("New UniFi client tracker %s (%s)", + client.hostname, client.mac) + + for device_id in controller.api.devices: + + if device_id in tracked: + LOGGER.debug("Updating UniFi tracked device %s (%s)", + tracked[device_id].entity_id, + tracked[device_id].device.mac) + tracked[device_id].async_schedule_update_ha_state() + continue + + device = controller.api.devices[device_id] + + tracked[device_id] = UniFiDeviceTracker(device, controller) + new_tracked.append(tracked[device_id]) + LOGGER.debug("New UniFi device tracker %s (%s)", + device.name, device.mac) if new_tracked: async_add_entities(new_tracked) class UniFiClientTracker(ScannerEntity): - """Representation of a network device.""" + """Representation of a network client.""" def __init__(self, client, controller): - """Set up tracked device.""" + """Set up tracked client.""" self.client = client self.controller = controller @@ -151,7 +168,7 @@ class UniFiClientTracker(ScannerEntity): @property def is_connected(self): - """Return true if the device is connected to the network.""" + """Return true if the client is connected to the network.""" detection_time = self.controller.unifi_config.get( CONF_DETECTION_TIME, DEFAULT_DETECTION_TIME) @@ -162,12 +179,12 @@ class UniFiClientTracker(ScannerEntity): @property def source_type(self): - """Return the source type of the device.""" + """Return the source type of the client.""" return SOURCE_TYPE_ROUTER @property def name(self) -> str: - """Return the name of the device.""" + """Return the name of the client.""" return self.client.name or self.client.hostname @property @@ -182,14 +199,14 @@ class UniFiClientTracker(ScannerEntity): @property def device_info(self): - """Return a device description for device registry.""" + """Return a client description for device registry.""" return { 'connections': {(CONNECTION_NETWORK_MAC, self.client.mac)} } @property def device_state_attributes(self): - """Return the device state attributes.""" + """Return the client state attributes.""" attributes = {} for variable in DEVICE_ATTRIBUTES: @@ -197,3 +214,71 @@ class UniFiClientTracker(ScannerEntity): attributes[variable] = self.client.raw[variable] return attributes + + +class UniFiDeviceTracker(ScannerEntity): + """Representation of a network infrastructure device.""" + + def __init__(self, device, controller): + """Set up tracked device.""" + self.device = device + self.controller = controller + + async def async_update(self): + """Synchronize state with controller.""" + await self.controller.request_update() + + @property + def is_connected(self): + """Return true if the device is connected to the network.""" + detection_time = self.controller.unifi_config.get( + CONF_DETECTION_TIME, DEFAULT_DETECTION_TIME) + + if (dt_util.utcnow() - dt_util.utc_from_timestamp(float( + self.device.last_seen))) < detection_time: + return True + return False + + @property + def source_type(self): + """Return the source type of the device.""" + return SOURCE_TYPE_ROUTER + + @property + def name(self) -> str: + """Return the name of the device.""" + return self.device.name + + @property + def unique_id(self) -> str: + """Return a unique identifier for this device.""" + return self.device.mac + + @property + def available(self) -> bool: + """Return if controller is available.""" + return self.controller.available + + @property + def device_info(self): + """Return a device description for device registry.""" + return { + 'connections': {(CONNECTION_NETWORK_MAC, self.device.mac)}, + 'manufacturer': ATTR_MANUFACTURER, + 'model': self.device.model, + 'name': self.device.name, + 'sw_version': self.device.version + } + + @property + def device_state_attributes(self): + """Return the device state attributes.""" + attributes = {} + + attributes['upgradable'] = self.device.upgradable + attributes['overheating'] = self.device.overheating + + if self.device.has_fan: + attributes['fan_level'] = self.device.fan_level + + return attributes diff --git a/homeassistant/components/unifi/manifest.json b/homeassistant/components/unifi/manifest.json index fff731abf4e..dc5e89c147e 100644 --- a/homeassistant/components/unifi/manifest.json +++ b/homeassistant/components/unifi/manifest.json @@ -4,7 +4,7 @@ "config_flow": true, "documentation": "https://www.home-assistant.io/components/unifi", "requirements": [ - "aiounifi==7" + "aiounifi==8" ], "dependencies": [], "codeowners": [ diff --git a/requirements_all.txt b/requirements_all.txt index 129b0d10173..86e75fce6a4 100644 --- a/requirements_all.txt +++ b/requirements_all.txt @@ -169,7 +169,7 @@ aiopvapi==1.6.14 aioswitcher==2019.4.26 # homeassistant.components.unifi -aiounifi==7 +aiounifi==8 # homeassistant.components.wwlln aiowwlln==1.0.0 diff --git a/requirements_test_all.txt b/requirements_test_all.txt index fbc12e7606e..2274bc65358 100644 --- a/requirements_test_all.txt +++ b/requirements_test_all.txt @@ -64,7 +64,7 @@ aionotion==1.1.0 aioswitcher==2019.4.26 # homeassistant.components.unifi -aiounifi==7 +aiounifi==8 # homeassistant.components.wwlln aiowwlln==1.0.0 diff --git a/tests/components/unifi/test_device_tracker.py b/tests/components/unifi/test_device_tracker.py index f8be08975fa..298649d0bfd 100644 --- a/tests/components/unifi/test_device_tracker.py +++ b/tests/components/unifi/test_device_tracker.py @@ -49,6 +49,22 @@ CLIENT_3 = { 'mac': '00:00:00:00:00:03', } +DEVICE_1 = { + 'board_rev': 3, + 'device_id': 'mock-id', + 'has_fan': True, + 'fan_level': 0, + 'ip': '10.0.1.1', + 'last_seen': 1562600145, + 'mac': '00:00:00:00:01:01', + 'model': 'US16P150', + 'name': 'device_1', + 'overheating': False, + 'type': 'usw', + 'upgradable': False, + 'version': '4.0.42.10433', +} + CONTROLLER_DATA = { CONF_HOST: 'mock-host', CONF_USERNAME: 'mock-user', @@ -137,32 +153,41 @@ async def test_tracked_devices(hass, mock_controller): """Test the update_items function with some clients.""" mock_controller.mock_client_responses.append( [CLIENT_1, CLIENT_2, CLIENT_3]) - mock_controller.mock_device_responses.append({}) + mock_controller.mock_device_responses.append([DEVICE_1]) mock_controller.unifi_config = {unifi_dt.CONF_SSID_FILTER: ['ssid']} await setup_controller(hass, mock_controller) assert len(mock_controller.mock_requests) == 2 - assert len(hass.states.async_all()) == 4 + assert len(hass.states.async_all()) == 5 - device_1 = hass.states.get('device_tracker.client_1') + client_1 = hass.states.get('device_tracker.client_1') + assert client_1 is not None + assert client_1.state == 'not_home' + + client_2 = hass.states.get('device_tracker.wired_client') + assert client_2 is not None + assert client_2.state == 'not_home' + + client_3 = hass.states.get('device_tracker.client_3') + assert client_3 is None + + device_1 = hass.states.get('device_tracker.device_1') assert device_1 is not None assert device_1.state == 'not_home' - device_2 = hass.states.get('device_tracker.wired_client') - assert device_2 is not None - assert device_2.state == 'not_home' - - device_3 = hass.states.get('device_tracker.client_3') - assert device_3 is None - - client_1 = copy(CLIENT_1) - client_1['last_seen'] = dt_util.as_timestamp(dt_util.utcnow()) - mock_controller.mock_client_responses.append([client_1]) - mock_controller.mock_device_responses.append({}) + client_1_copy = copy(CLIENT_1) + client_1_copy['last_seen'] = dt_util.as_timestamp(dt_util.utcnow()) + device_1_copy = copy(DEVICE_1) + device_1_copy['last_seen'] = dt_util.as_timestamp(dt_util.utcnow()) + mock_controller.mock_client_responses.append([client_1_copy]) + mock_controller.mock_device_responses.append([device_1_copy]) await mock_controller.async_update() await hass.async_block_till_done() - device_1 = hass.states.get('device_tracker.client_1') + client_1 = hass.states.get('device_tracker.client_1') + assert client_1.state == 'home' + + device_1 = hass.states.get('device_tracker.device_1') assert device_1.state == 'home' diff --git a/tests/components/unifi/test_init.py b/tests/components/unifi/test_init.py index 2d9ea143e76..063975406ab 100644 --- a/tests/components/unifi/test_init.py +++ b/tests/components/unifi/test_init.py @@ -77,7 +77,7 @@ async def test_successful_config_entry(hass): 'connections': { ('mac', '00:11:22:33:44:55') }, - 'manufacturer': 'Ubiquiti', + 'manufacturer': unifi.ATTR_MANUFACTURER, 'model': "UniFi Controller", 'name': "UniFi Controller", }