Changes after review by @MartinHjelmare

This commit is contained in:
Steven Looman 2018-10-01 18:25:54 +02:00
parent f511920a04
commit d732f8eca2
9 changed files with 147 additions and 115 deletions

View File

@ -8,8 +8,10 @@ https://home-assistant.io/components/sensor.upnp/
from datetime import datetime
import logging
from homeassistant.components.upnp import DOMAIN
from homeassistant.core import callback
from homeassistant.helpers.dispatcher import async_dispatcher_connect
from homeassistant.helpers.entity import Entity
from homeassistant.components.upnp.const import DOMAIN as DATA_UPNP
_LOGGER = logging.getLogger(__name__)
@ -45,37 +47,65 @@ OUT = 'sent'
KBYTE = 1024
async def async_setup_platform(hass, config, async_add_devices,
async def async_setup_platform(hass, config, async_add_entities,
discovery_info=None):
"""Set up the UPnP/IGD sensors."""
if discovery_info is None:
return
udn = discovery_info['udn']
device = hass.data[DOMAIN]['devices'][udn]
# raw sensors + per-second sensors
sensors = [
RawUPnPIGDSensor(device, name, sensor_type)
for name, sensor_type in SENSOR_TYPES.items()
]
sensors += [
KBytePerSecondUPnPIGDSensor(device, IN),
KBytePerSecondUPnPIGDSensor(device, OUT),
PacketsPerSecondUPnPIGDSensor(device, IN),
PacketsPerSecondUPnPIGDSensor(device, OUT),
]
hass.data[DOMAIN]['sensors'][udn] = sensors
async_add_devices(sensors, True)
return True
"""Old way of setting up UPnP/IGD sensors."""
_LOGGER.debug('async_setup_platform: config: %s, discovery: %s',
config, discovery_info)
class RawUPnPIGDSensor(Entity):
async def async_setup_entry(hass, config_entry, async_add_entities):
"""Set up the UPnP/IGD sensor."""
@callback
def async_add_sensor(device):
"""Add sensors from UPnP/IGD device."""
# raw sensors + per-second sensors
sensors = [
RawUPnPIGDSensor(device, name, sensor_type)
for name, sensor_type in SENSOR_TYPES.items()
]
sensors += [
KBytePerSecondUPnPIGDSensor(device, IN),
KBytePerSecondUPnPIGDSensor(device, OUT),
PacketsPerSecondUPnPIGDSensor(device, IN),
PacketsPerSecondUPnPIGDSensor(device, OUT),
]
async_add_entities(sensors, True)
data = config_entry.data
udn = data['udn']
device = hass.data[DATA_UPNP]['devices'][udn]
async_add_sensor(device)
class UpnpSensor(Entity):
"""Base class for UPnP/IGD sensors."""
def __init__(self, device):
"""Initialize the base sensor."""
self._device = device
async def async_added_to_hass(self):
"""Subscribe to sensors events."""
async_dispatcher_connect(self.hass,
'upnp_remove_sensor',
self._upnp_remove_sensor)
def _upnp_remove_sensor(self, device):
"""Remove sensor."""
if self._device != device:
# not for us
return
self.hass.async_create_task(self.async_remove())
class RawUPnPIGDSensor(UpnpSensor):
"""Representation of a UPnP/IGD sensor."""
def __init__(self, device, sensor_type_name, sensor_type):
"""Initialize the UPnP/IGD sensor."""
self._device = device
super().__init__(device)
self._type_name = sensor_type_name
self._type = sensor_type
self._name = '{} {}'.format(device.name, sensor_type['name'])
@ -94,7 +124,7 @@ class RawUPnPIGDSensor(Entity):
@property
def state(self) -> str:
"""Return the state of the device."""
return self._state
return format(self._state, 'd')
@property
def icon(self) -> str:
@ -118,12 +148,12 @@ class RawUPnPIGDSensor(Entity):
self._state = await self._device.async_get_total_packets_sent()
class PerSecondUPnPIGDSensor(Entity):
class PerSecondUPnPIGDSensor(UpnpSensor):
"""Abstract representation of a X Sent/Received per second sensor."""
def __init__(self, device, direction):
"""Initializer."""
self._device = device
super().__init__(device)
self._direction = direction
self._state = None
@ -205,7 +235,7 @@ class KBytePerSecondUPnPIGDSensor(PerSecondUPnPIGDSensor):
return await self._device.async_get_total_bytes_sent()
@property
def state(self):
def state(self) -> str:
"""Return the state of the device."""
if self._state is None:
return None
@ -229,7 +259,7 @@ class PacketsPerSecondUPnPIGDSensor(PerSecondUPnPIGDSensor):
return await self._device.async_get_total_packets_sent()
@property
def state(self):
def state(self) -> str:
"""Return the state of the device."""
if self._state is None:
return None

View File

@ -9,8 +9,8 @@
"title": "Configuration options for the UPnP/IGD",
"data":{
"igd": "UPnP/IGD",
"sensors": "Add traffic sensors",
"port_forward": "Enable port forward for Home Assistant"
"enable_sensors": "Add traffic sensors",
"enable_port_mapping": "Enable port mapping for Home Assistant"
}
}
},
@ -19,7 +19,7 @@
"abort": {
"no_devices_discovered": "No UPnP/IGDs discovered",
"already_configured": "UPnP/IGD is already configured",
"no_sensors_or_port_forward": "Enable at least sensors or port forward"
"no_sensors_or_port_mapping": "Enable at least sensors or port mapping"
}
}
}

View File

@ -9,8 +9,8 @@
"title": "Extra configuratie options voor UPnP/IGD",
"data":{
"igd": "UPnP/IGD",
"sensors": "Verkeer sensors toevoegen",
"port_forward": "Maak port forward voor Home Assistant"
"enable_sensors": "Verkeer sensors toevoegen",
"enable_port_mapping": "Maak port mapping voor Home Assistant"
}
}
},
@ -19,7 +19,7 @@
"abort": {
"no_devices_discovered": "Geen UPnP/IGDs gevonden",
"already_configured": "UPnP/IGD is reeds geconfigureerd",
"no_sensors_or_port_forward": "Kies ten minste sensors of port forward"
"no_sensors_or_port_mapping": "Kies ten minste sensors of port mapping"
}
}
}

View File

@ -4,6 +4,7 @@ Will open a port in your router for Home Assistant and provide statistics.
For more details about this component, please refer to the documentation at
https://home-assistant.io/components/upnp/
"""
# pylint: disable=invalid-name
import asyncio
from ipaddress import ip_address
@ -12,9 +13,8 @@ import voluptuous as vol
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import EVENT_HOMEASSISTANT_STOP
from homeassistant.exceptions import PlatformNotReady
from homeassistant.helpers import config_validation as cv
from homeassistant.helpers import discovery
from homeassistant.helpers import dispatcher
from homeassistant.helpers.typing import ConfigType
from homeassistant.helpers.typing import HomeAssistantType
from homeassistant.components.discovery import DOMAIN as DISCOVERY_DOMAIN
@ -79,16 +79,16 @@ async def async_setup(hass: HomeAssistantType, config: ConfigType):
hass.data[DOMAIN]['local_ip'] = upnp_config[CONF_LOCAL_IP]
# determine ports
ports = {CONF_HASS: CONF_HASS} # default, port_forward disabled by default
ports = {CONF_HASS: CONF_HASS} # default, port_mapping disabled by default
if CONF_PORTS in upnp_config:
# copy from config
ports = upnp_config[CONF_PORTS]
hass.data[DOMAIN]['auto_config'] = {
'active': True,
'port_forward': upnp_config[CONF_ENABLE_PORT_MAPPING],
'enable_sensors': upnp_config[CONF_ENABLE_SENSORS],
'enable_port_mapping': upnp_config[CONF_ENABLE_PORT_MAPPING],
'ports': ports,
'sensors': upnp_config[CONF_ENABLE_SENSORS],
}
return True
@ -98,7 +98,6 @@ async def async_setup(hass: HomeAssistantType, config: ConfigType):
async def async_setup_entry(hass: HomeAssistantType,
config_entry: ConfigEntry):
"""Set up a bridge from a config entry."""
_LOGGER.debug('async_setup_entry: %s', config_entry.data)
ensure_domain_data(hass)
data = config_entry.data
@ -107,7 +106,8 @@ async def async_setup_entry(hass: HomeAssistantType,
try:
device = await Device.async_create_device(hass, ssdp_description)
except (asyncio.TimeoutError, aiohttp.ClientError):
raise PlatformNotReady()
_LOGGER.error('Unable to create upnp-device')
return
hass.data[DOMAIN]['devices'][device.udn] = device
@ -124,12 +124,10 @@ async def async_setup_entry(hass: HomeAssistantType,
# sensors
if data.get(CONF_ENABLE_SENSORS):
_LOGGER.debug('Enabling sensors')
discovery_info = {
'udn': device.udn,
}
hass_config = config_entry.data
hass.async_create_task(discovery.async_load_platform(
hass, 'sensor', DOMAIN, discovery_info, hass_config))
# register sensor setup handlers
hass.async_create_task(hass.config_entries.async_forward_entry_setup(
config_entry, 'sensor'))
async def unload_entry(event):
"""Unload entry on quit."""
@ -142,7 +140,6 @@ async def async_setup_entry(hass: HomeAssistantType,
async def async_unload_entry(hass: HomeAssistantType,
config_entry: ConfigEntry):
"""Unload a config entry."""
_LOGGER.debug('async_unload_entry: %s', config_entry.data)
data = config_entry.data
udn = data[CONF_UDN]
@ -156,9 +153,9 @@ async def async_unload_entry(hass: HomeAssistantType,
await device.async_delete_port_mappings()
# sensors
for sensor in hass.data[DOMAIN]['sensors'].get(udn, []):
_LOGGER.debug('Deleting sensor: %s', sensor)
await sensor.async_remove()
if data.get(CONF_ENABLE_SENSORS):
_LOGGER.debug('Deleting sensors')
dispatcher.async_dispatcher_send(hass, 'upnp_remove_sensor', device)
# clear stored device
del hass.data[DOMAIN]['devices'][udn]

View File

@ -1,4 +1,6 @@
"""Config flow for UPNP."""
from collections import OrderedDict
import voluptuous as vol
from homeassistant import config_entries
@ -15,12 +17,11 @@ def ensure_domain_data(hass):
"""Ensure hass.data is filled properly."""
hass.data[DOMAIN] = hass.data.get(DOMAIN, {})
hass.data[DOMAIN]['devices'] = hass.data[DOMAIN].get('devices', {})
hass.data[DOMAIN]['sensors'] = hass.data[DOMAIN].get('sensors', {})
hass.data[DOMAIN]['discovered'] = hass.data[DOMAIN].get('discovered', {})
hass.data[DOMAIN]['auto_config'] = hass.data[DOMAIN].get('auto_config', {
'active': False,
'port_forward': False,
'sensors': False,
'enable_sensors': False,
'enable_port_mapping': False,
'ports': {'hass': 'hass'},
})
@ -30,6 +31,7 @@ class UpnpFlowHandler(data_entry_flow.FlowHandler):
"""Handle a Hue config flow."""
VERSION = 1
CONNECTION_CLASS = config_entries.CONN_CLASS_LOCAL_POLL
@property
def _configured_upnp_igds(self):
@ -79,8 +81,8 @@ class UpnpFlowHandler(data_entry_flow.FlowHandler):
if auto_config['active']:
import_info = {
'name': discovery_info['friendly_name'],
'sensors': auto_config['sensors'],
'port_forward': auto_config['port_forward'],
'enable_sensors': auto_config['enable_sensors'],
'enable_port_mapping': auto_config['enable_port_mapping'],
}
return await self._async_save_entry(import_info)
@ -94,8 +96,9 @@ class UpnpFlowHandler(data_entry_flow.FlowHandler):
# if user input given, handle it
user_input = user_input or {}
if 'name' in user_input:
if not user_input['sensors'] and not user_input['port_forward']:
return self.async_abort(reason='no_sensors_or_port_forward')
if not user_input['enable_sensors'] and \
not user_input['enable_port_mapping']:
return self.async_abort(reason='no_sensors_or_port_mapping')
# ensure not already configured
configured_names = [
@ -119,12 +122,13 @@ class UpnpFlowHandler(data_entry_flow.FlowHandler):
return self.async_show_form(
step_id='user',
data_schema=vol.Schema({
vol.Required('name'): vol.In(names),
vol.Optional('sensors', default=False): bool,
vol.Optional('port_forward', default=False): bool,
})
)
data_schema=vol.Schema(
OrderedDict([
(vol.Required('name'), vol.In(names)),
(vol.Optional('enable_sensors', default=False), bool),
(vol.Optional('enable_port_mapping', default=False), bool),
])
))
async def async_step_import(self, import_info):
"""Import a new UPnP/IGD as a config entry."""
@ -150,7 +154,7 @@ class UpnpFlowHandler(data_entry_flow.FlowHandler):
data={
CONF_SSDP_DESCRIPTION: discovery_info['ssdp_description'],
CONF_UDN: discovery_info['udn'],
CONF_ENABLE_SENSORS: import_info['sensors'],
CONF_ENABLE_PORT_MAPPING: import_info['port_forward'],
CONF_ENABLE_SENSORS: import_info['enable_sensors'],
CONF_ENABLE_PORT_MAPPING: import_info['enable_port_mapping'],
},
)

View File

@ -39,7 +39,7 @@ class Device:
from async_upnp_client.igd import IgdDevice
igd_device = IgdDevice(upnp_device, None)
return Device(igd_device)
return cls(igd_device)
@property
def udn(self):
@ -102,6 +102,7 @@ class Device:
async def _async_delete_port_mapping(self, external_port):
"""Remove a port mapping."""
from async_upnp_client import UpnpError
_LOGGER.info('Deleting port mapping %s (TCP)', external_port)
try:
await self._igd_device.async_delete_port_mapping(
remote_host=None,

View File

@ -9,8 +9,8 @@
"title": "Configuration options for the UPnP/IGD",
"data":{
"igd": "UPnP/IGD",
"sensors": "Add traffic sensors",
"port_forward": "Enable port forward for Home Assistant"
"enable_sensors": "Add traffic sensors",
"enable_port_mapping": "Enable port mapping for Home Assistant"
}
}
},
@ -19,7 +19,7 @@
"abort": {
"no_devices_discovered": "No UPnP/IGDs discovered",
"already_configured": "UPnP/IGD is already configured",
"no_sensors_or_port_forward": "Enable at least sensors or port forward"
"no_sensors_or_port_mapping": "Enable at least sensors or port mapping"
}
}
}

View File

@ -44,15 +44,15 @@ async def test_flow_already_configured(hass):
result = await flow.async_step_user({
'name': '192.168.1.1 (Test device)',
'sensors': True,
'port_forward': False,
'enable_sensors': True,
'enable_port_mapping': False,
})
assert result['type'] == 'abort'
assert result['reason'] == 'already_configured'
async def test_flow_no_sensors_no_port_forward(hass):
"""Test single device, no sensors, no port_forward."""
async def test_flow_no_sensors_no_port_mapping(hass):
"""Test single device, no sensors, no port_mapping."""
flow = upnp_config_flow.UpnpFlowHandler()
flow.hass = hass
@ -76,11 +76,11 @@ async def test_flow_no_sensors_no_port_forward(hass):
result = await flow.async_step_user({
'name': '192.168.1.1 (Test device)',
'sensors': False,
'port_forward': False,
'enable_sensors': False,
'enable_port_mapping': False,
})
assert result['type'] == 'abort'
assert result['reason'] == 'no_sensors_or_port_forward'
assert result['reason'] == 'no_sensors_or_port_mapping'
async def test_flow_discovered_form(hass):
@ -106,7 +106,7 @@ async def test_flow_discovered_form(hass):
async def test_flow_two_discovered_form(hass):
"""Test single device discovered, show form flow."""
"""Test two devices discovered, show form flow with two devices."""
flow = upnp_config_flow.UpnpFlowHandler()
flow.hass = hass
@ -133,13 +133,13 @@ async def test_flow_two_discovered_form(hass):
assert result['step_id'] == 'user'
assert result['data_schema']({
'name': '192.168.1.1 (Test device)',
'sensors': True,
'port_forward': False,
'enable_sensors': True,
'enable_port_mapping': False,
})
assert result['data_schema']({
'name': '192.168.1.1 (Test device)',
'sensors': True,
'port_forward': False,
'name': '192.168.2.1 (Test device)',
'enable_sensors': True,
'enable_port_mapping': False,
})
@ -163,15 +163,15 @@ async def test_config_entry_created(hass):
result = await flow.async_step_user({
'name': '192.168.1.1 (Test device)',
'sensors': True,
'port_forward': False,
'enable_sensors': True,
'enable_port_mapping': False,
})
assert result['type'] == 'create_entry'
assert result['data'] == {
'port_forward': False,
'sensors': True,
'ssdp_description': 'http://192.168.1.1/desc.xml',
'udn': 'uuid:device_1',
'port_mapping': False,
'sensors': True,
}
assert result['title'] == 'Test device 1'
@ -185,8 +185,8 @@ async def test_flow_discovery_auto_config_sensors(hass):
hass.data[upnp.DOMAIN] = {
'auto_config': {
'active': True,
'port_forward': False,
'sensors': True,
'enable_port_mapping': False,
'enable_sensors': True,
},
}
@ -200,25 +200,25 @@ async def test_flow_discovery_auto_config_sensors(hass):
assert result['type'] == 'create_entry'
assert result['data'] == {
'port_forward': False,
'sensors': True,
'ssdp_description': 'http://192.168.1.1/desc.xml',
'udn': 'uuid:device_1',
'sensors': True,
'port_mapping': False,
}
assert result['title'] == 'Test device 1'
async def test_flow_discovery_auto_config_sensors_port_forward(hass):
"""Test creation of device with auto_config, with port forward."""
async def test_flow_discovery_auto_config_sensors_port_mapping(hass):
"""Test creation of device with auto_config, with port mapping."""
flow = upnp_config_flow.UpnpFlowHandler()
flow.hass = hass
# auto_config active, with port_forward
# auto_config active, with port_mapping
hass.data[upnp.DOMAIN] = {
'auto_config': {
'active': True,
'port_forward': True,
'sensors': True,
'enable_port_mapping': True,
'enable_sensors': True,
},
}
@ -232,9 +232,9 @@ async def test_flow_discovery_auto_config_sensors_port_forward(hass):
assert result['type'] == 'create_entry'
assert result['data'] == {
'port_forward': True,
'sensors': True,
'ssdp_description': 'http://192.168.1.1/desc.xml',
'udn': 'uuid:device_1',
'ssdp_description': 'http://192.168.1.1/desc.xml',
'sensors': True,
'port_mapping': True,
}
assert result['title'] == 'Test device 1'

View File

@ -53,9 +53,9 @@ async def test_async_setup_no_auto_config(hass):
assert hass.data[upnp.DOMAIN]['auto_config'] == {
'active': False,
'port_forward': False,
'enable_sensors': False,
'enable_port_mapping': False,
'ports': {'hass': 'hass'},
'sensors': False,
}
@ -66,27 +66,27 @@ async def test_async_setup_auto_config(hass):
assert hass.data[upnp.DOMAIN]['auto_config'] == {
'active': True,
'port_forward': False,
'enable_sensors': True,
'enable_port_mapping': False,
'ports': {'hass': 'hass'},
'sensors': True,
}
async def test_async_setup_auto_config_port_forward(hass):
async def test_async_setup_auto_config_port_mapping(hass):
"""Test async_setup."""
# setup component, enable auto_config
await async_setup_component(hass, 'upnp', {
'upnp': {
'port_forward': True,
'port_mapping': True,
'ports': {'hass': 'hass'},
},
'discovery': {}})
assert hass.data[upnp.DOMAIN]['auto_config'] == {
'active': True,
'port_forward': True,
'enable_sensors': True,
'enable_port_mapping': True,
'ports': {'hass': 'hass'},
'sensors': True,
}
@ -99,9 +99,9 @@ async def test_async_setup_auto_config_no_sensors(hass):
assert hass.data[upnp.DOMAIN]['auto_config'] == {
'active': True,
'port_forward': False,
'enable_sensors': False,
'enable_port_mapping': False,
'ports': {'hass': 'hass'},
'sensors': False,
}
@ -112,7 +112,7 @@ async def test_async_setup_entry_default(hass):
'ssdp_description': 'http://192.168.1.1/desc.xml',
'udn': udn,
'sensors': True,
'port_forward': False,
'port_mapping': False,
})
# ensure hass.http is available
@ -144,20 +144,20 @@ async def test_async_setup_entry_default(hass):
assert len(mock_device.async_delete_port_mappings.mock_calls) == 0
async def test_async_setup_entry_port_forward(hass):
async def test_async_setup_entry_port_mapping(hass):
"""Test async_setup_entry."""
udn = 'uuid:device_1'
entry = MockConfigEntry(domain=upnp.DOMAIN, data={
'ssdp_description': 'http://192.168.1.1/desc.xml',
'udn': udn,
'sensors': False,
'port_forward': True,
'port_mapping': True,
})
# ensure hass.http is available
await async_setup_component(hass, 'upnp', {
'upnp': {
'port_forward': True,
'port_mapping': True,
'ports': {'hass': 'hass'},
},
'discovery': {},