Working on igd

CODEOWNERS

@@ -80,7 +80,6 @@ homeassistant/components/sensor/sma.py @kellerza
 homeassistant/components/sensor/sql.py @dgomes
 homeassistant/components/sensor/sytadin.py @gautric
 homeassistant/components/sensor/tibber.py @danielhiversen
-homeassistant/components/sensor/upnp.py @dgomes
 homeassistant/components/sensor/waqi.py @andrey-git
 homeassistant/components/switch/tplink.py @rytilahti
 homeassistant/components/vacuum/roomba.py @pschmitt

homeassistant/components/igd/.translations/en.json

@@ -10,7 +10,7 @@
                 "data":{
                     "igd": "IGD",
                     "sensors": "Add traffic sensors",
-                    "port_forward": "Enable port forward for Home Assistant<br>Only enable this when your Home Assistant is password protected!"
+                    "port_forward": "Enable port forward for Home Assistant"
                 }
             }
         },
@@ -20,8 +20,7 @@
             "no_devices_discovered": "No IGDs discovered",
             "already_configured": "IGD is already configured",
             "no_sensors_or_port_forward": "Enable at least sensors or Port forward",
-            "no_igds": "No IGDs discovered",
+            "no_igds": "No IGDs discovered"
-            "todo": "TODO"
         }
     }
 }

homeassistant/components/igd/.translations/nl.json

@@ -10,7 +10,7 @@
                 "data":{
                     "igd": "IGD",
                     "sensors": "Verkeer sensors toevoegen",
-                    "port_forward": "Maak port-forward voor Home Assistant<br>Zet dit alleen aan wanneer uw Home Assistant een wachtwoord heeft!"
+                    "port_forward": "Maak port-forward voor Home Assistant"
                 }
             }
         },
@@ -20,8 +20,7 @@
             "no_devices_discovered": "Geen IGDs gevonden",
             "already_configured": "IGD is reeds geconfigureerd",
             "no_sensors_or_port_forward": "Kies ten minste sensors of Port forward",
-            "no_igds": "Geen IGDs gevonden",
+            "no_igds": "Geen IGDs gevonden"
-            "todo": "TODO"
         }
     }
 }

homeassistant/components/igd/__init__.py

@@ -2,10 +2,9 @@
 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/
+https://home-assistant.io/components/igd/
 """
 
-
 import asyncio
 from ipaddress import IPv4Address
 from ipaddress import ip_address
@@ -22,19 +21,22 @@ from homeassistant.helpers.aiohttp_client import async_get_clientsession
 from homeassistant.helpers.typing import ConfigType, HomeAssistantType
 from homeassistant.util import get_local_ip
 from homeassistant.components.discovery import DOMAIN as DISCOVERY_DOMAIN
+import homeassistant.components.igd.config_flow  # noqa: 401
 
-from .const import CONF_ENABLE_PORT_MAPPING, CONF_ENABLE_SENSORS
+from .const import (
+    CONF_ENABLE_PORT_MAPPING, CONF_ENABLE_SENSORS,
+    CONF_UDN, CONF_SSDP_DESCRIPTION
+)
 from .const import DOMAIN
 from .const import LOGGER as _LOGGER
-import homeassistant.components.igd.config_flow  # register the handler
+from .const import ensure_domain_data
 
 
 REQUIREMENTS = ['async-upnp-client==0.12.4']
-DEPENDENCIES = ['http']  # ,'discovery']
+DEPENDENCIES = ['http']
 
 CONF_LOCAL_IP = 'local_ip'
 CONF_PORTS = 'ports'
-CONF_HASS = 'hass'
 
 NOTIFICATION_ID = 'igd_notification'
 NOTIFICATION_TITLE = 'UPnP/IGD Setup'
@@ -72,16 +74,15 @@ async def _async_create_igd_device(hass: HomeAssistantType,
 
 def _store_device(hass: HomeAssistantType, udn, igd_device):
     """Store an igd_device by udn."""
-    hass.data[DOMAIN] = hass.data.get(DOMAIN, {})
+    if igd_device is not None:
-    hass.data[DOMAIN]['devices'] = hass.data[DOMAIN].get('devices', {})
         hass.data[DOMAIN]['devices'][udn] = igd_device
+    elif udn in hass.data[DOMAIN]['devices']:
+        del hass.data[DOMAIN]['devices'][udn]
 
 
 def _get_device(hass: HomeAssistantType, udn):
     """Get an igd_device by udn."""
-    hass.data[DOMAIN] = hass.data.get(DOMAIN, {})
+    return hass.data[DOMAIN]['devices'].get(udn)
-    hass.data[DOMAIN]['devices'] = hass.data[DOMAIN].get('devices', {})
-    return hass.data[DOMAIN]['devices'][udn]
 
 
 async def _async_add_port_mapping(hass: HomeAssistantType,
@@ -123,14 +124,7 @@ async def _async_delete_port_mapping(hass: HomeAssistantType, igd_device):
 # config
 async def async_setup(hass: HomeAssistantType, config: ConfigType):
     """Register a port mapping for Home Assistant via UPnP."""
-    # defaults
+    ensure_domain_data(hass)
-    hass.data[DOMAIN] = hass.data.get(DOMAIN, {})
-    if 'auto_config' not in hass.data[DOMAIN]:
-        hass.data[DOMAIN]['auto_config'] = {
-            'active': False,
-            'port_forward': False,
-            'sensors': False,
-        }
 
     # ensure sane config
     if DOMAIN not in config:
@@ -158,12 +152,11 @@ 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, %s', config_entry, config_entry.entry_id)
+    ensure_domain_data(hass)
-
     data = config_entry.data
-    ssdp_description = data['ssdp_description']
 
     # build IGD device
+    ssdp_description = data[CONF_SSDP_DESCRIPTION]
     try:
         igd_device = await _async_create_igd_device(hass, ssdp_description)
     except (asyncio.TimeoutError, aiohttp.ClientError):
@@ -179,11 +172,10 @@ async def async_setup_entry(hass: HomeAssistantType,
     # sensors
     if data.get(CONF_ENABLE_SENSORS):
         discovery_info = {
-            'udn': data['udn'],
+            'udn': data[CONF_UDN],
         }
         hass_config = config_entry.data
-        hass.async_create_task(
+        hass.async_create_task(discovery.async_load_platform(
-            discovery.async_load_platform(
             hass, 'sensor', DOMAIN, discovery_info, hass_config))
 
     async def unload_entry(event):
@@ -197,12 +189,8 @@ 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, entry_id: %s, data: %s', config_entry, config_entry.entry_id, config_entry.data)
-    for entry in hass.config_entries._entries:
-        _LOGGER.debug('%s: %s: %s', entry, entry.entry_id, entry.data)
-
     data = config_entry.data
-    udn = data['udn']
+    udn = data[CONF_UDN]
 
     igd_device = _get_device(hass, udn)
     if igd_device is None:
@@ -213,14 +201,10 @@ async def async_unload_entry(hass: HomeAssistantType,
         await _async_delete_port_mapping(hass, igd_device)
 
     # sensors
-    if data.get(CONF_ENABLE_SENSORS):
+    for sensor in hass.data[DOMAIN]['sensors'].get(udn, []):
-        # XXX TODO: remove sensors
+        await sensor.async_remove()
-        pass
 
     # clear stored device
     _store_device(hass, udn, None)
 
-    # XXX TODO: remove config entry
-    #await hass.config_entries.async_remove(config_entry.entry_id)
-
     return True

homeassistant/components/igd/config_flow.py

@@ -1,21 +1,15 @@
 """Config flow for IGD."""
 import voluptuous as vol
 
-from homeassistant import config_entries, data_entry_flow
+from homeassistant import config_entries
-from homeassistant.core import callback
+from homeassistant import data_entry_flow
 
-from .const import CONF_ENABLE_PORT_MAPPING, CONF_ENABLE_SENSORS
+from .const import (
+    CONF_ENABLE_PORT_MAPPING, CONF_ENABLE_SENSORS,
+    CONF_SSDP_DESCRIPTION, CONF_UDN
+)
 from .const import DOMAIN
-from .const import LOGGER as _LOGGER
+from .const import ensure_domain_data
-
-
-@callback
-def configured_udns(hass):
-    """Get all configured UDNs."""
-    return [
-        entry.data['udn']
-        for entry in hass.config_entries.async_entries(DOMAIN)
-    ]
 
 
 @config_entries.HANDLERS.register(DOMAIN)
@@ -24,29 +18,29 @@ class IgdFlowHandler(data_entry_flow.FlowHandler):
 
     VERSION = 1
 
-    def __init__(self):
+    @property
-        """Initializer."""
+    def _configured_igds(self):
-        pass
+        """Get all configured IGDs."""
+        return {
+            entry.data[CONF_UDN]: {
+                'udn': entry.data[CONF_UDN],
+            }
+            for entry in self.hass.config_entries.async_entries(DOMAIN)
+        }
 
     @property
-    def _discovereds(self):
+    def _discovered_igds(self):
         """Get all discovered entries."""
-        return self.hass.data.get(DOMAIN, {}).get('discovered', {})
+        return self.hass.data[DOMAIN]['discovered']
 
     def _store_discovery_info(self, discovery_info):
         """Add discovery info."""
         udn = discovery_info['udn']
-        self.hass.data[DOMAIN] = self.hass.data.get(DOMAIN, {})
-        self.hass.data[DOMAIN]['discovered'] = \
-            self.hass.data[DOMAIN].get('discovered', {})
         self.hass.data[DOMAIN]['discovered'][udn] = discovery_info
 
     def _auto_config_settings(self):
         """Check if auto_config has been enabled."""
-        self.hass.data[DOMAIN] = self.hass.data.get(DOMAIN, {})
+        return self.hass.data[DOMAIN]['auto_config']
-        return self.hass.data[DOMAIN].get('auto_config', {
-            'active': False,
-        })
 
     async def async_step_discovery(self, discovery_info):
         """
@@ -55,16 +49,18 @@ class IgdFlowHandler(data_entry_flow.FlowHandler):
         This flow is triggered by the discovery component. It will check if the
         host is already configured and delegate to the import step if not.
         """
-        # ensure not already discovered/configured
+        ensure_domain_data(self.hass)
-        udn = discovery_info['udn']
-        if udn in configured_udns(self.hass):
-            return self.async_abort(reason='already_configured')
 
         # store discovered device
         discovery_info['friendly_name'] = \
             '{} ({})'.format(discovery_info['host'], discovery_info['name'])
         self._store_discovery_info(discovery_info)
 
+        # ensure not already discovered/configured
+        udn = discovery_info['udn']
+        if udn in self._configured_igds:
+            return self.async_abort(reason='already_configured')
+
         # auto config?
         auto_config = self._auto_config_settings()
         if auto_config['active']:
@@ -86,14 +82,13 @@ class IgdFlowHandler(data_entry_flow.FlowHandler):
             if not user_input['sensors'] and not user_input['port_forward']:
                 return self.async_abort(reason='no_sensors_or_port_forward')
 
-            # ensure nto already configured
+            # ensure not already configured
-            configured_igds = [
+            configured_names = [
                 entry['friendly_name']
-                for entry in self._discovereds.values()
+                for udn, entry in self._discovered_igds.items()
-                if entry['udn'] in configured_udns(self.hass)
+                if udn in self._configured_igds
             ]
-            _LOGGER.debug('Configured IGDs: %s', configured_igds)
+            if user_input['name'] in configured_names:
-            if user_input['name'] in configured_igds:
                 return self.async_abort(reason='already_configured')
 
             return await self._async_save_entry(user_input)
@@ -101,8 +96,8 @@ class IgdFlowHandler(data_entry_flow.FlowHandler):
         # let user choose from all discovered, non-configured, IGDs
         names = [
             entry['friendly_name']
-            for entry in self._discovereds.values()
+            for udn, entry in self._discovered_igds.items()
-            if entry['udn'] not in configured_udns(self.hass)
+            if udn not in self._configured_igds
         ]
         if not names:
             return self.async_abort(reason='no_devices_discovered')
@@ -125,7 +120,7 @@ class IgdFlowHandler(data_entry_flow.FlowHandler):
         # ensure we know the host
         name = import_info['name']
         discovery_infos = [info
-                           for info in self._discovereds.values()
+                           for info in self._discovered_igds.values()
                            if info['friendly_name'] == name]
         if not discovery_infos:
             return self.async_abort(reason='host_not_found')
@@ -134,8 +129,8 @@ class IgdFlowHandler(data_entry_flow.FlowHandler):
         return self.async_create_entry(
             title=discovery_info['name'],
             data={
-                'ssdp_description': discovery_info['ssdp_description'],
+                CONF_SSDP_DESCRIPTION: discovery_info['ssdp_description'],
-                'udn': discovery_info['udn'],
+                CONF_UDN: discovery_info['udn'],
                 CONF_ENABLE_SENSORS: import_info['sensors'],
                 CONF_ENABLE_PORT_MAPPING: import_info['port_forward'],
             },

homeassistant/components/igd/const.py

@@ -1,7 +1,23 @@
 """Constants for the IGD component."""
 import logging
 
+
 DOMAIN = 'igd'
 LOGGER = logging.getLogger('homeassistant.components.igd')
 CONF_ENABLE_PORT_MAPPING = 'port_forward'
 CONF_ENABLE_SENSORS = 'sensors'
+CONF_UDN = 'udn'
+CONF_SSDP_DESCRIPTION = 'ssdp_description'
+
+
+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,
+    })

homeassistant/components/igd/strings.json

@@ -10,7 +10,7 @@
                 "data":{
                     "igd": "IGD",
                     "sensors": "Add traffic sensors",
-                    "port_forward": "Enable port forward for Home Assistant<br>Only enable this when your Home Assistant is password protected!"
+                    "port_forward": "Enable port forward for Home Assistant"
                 }
             }
         },
@@ -20,8 +20,7 @@
             "no_devices_discovered": "No IGDs discovered",
             "already_configured": "IGD is already configured",
             "no_sensors_or_port_forward": "Enable at least sensors or Port forward",
-            "no_igds": "No IGDs discovered",
+            "no_igds": "No IGDs discovered"
-            "todo": "TODO"
         }
     }
 }

homeassistant/components/sensor/igd.py

@@ -32,11 +32,11 @@ SENSOR_TYPES = {
     },
     PACKETS_RECEIVED: {
         'name': 'packets received',
-        'unit': '#',
+        'unit': 'packets',
     },
     PACKETS_SENT: {
         'name': 'packets sent',
-        'unit': '#',
+        'unit': 'packets',
     },
 }
 
@@ -54,22 +54,20 @@ async def async_setup_platform(hass, config, async_add_devices,
     udn = discovery_info['udn']
     igd_device = hass.data[DOMAIN]['devices'][udn]
 
-    # raw sensors
+    # raw sensors + per-second sensors
-    async_add_devices([
+    sensors = [
         RawIGDSensor(igd_device, name, sensor_type)
-        for name, sensor_type in SENSOR_TYPES.items()],
+        for name, sensor_type in SENSOR_TYPES.items()
-        True
+    ]
-    )
+    sensors += [
-
-    # current traffic reporting
-    async_add_devices(
-        [
         KBytePerSecondIGDSensor(igd_device, IN),
         KBytePerSecondIGDSensor(igd_device, OUT),
         PacketsPerSecondIGDSensor(igd_device, IN),
         PacketsPerSecondIGDSensor(igd_device, OUT),
-        ], True
+    ]
-    )
+    hass.data[DOMAIN]['sensors'][udn] = sensors
+    async_add_devices(sensors, True)
+    return True
 
 
 class RawIGDSensor(Entity):
@@ -80,7 +78,7 @@ class RawIGDSensor(Entity):
         self._device = device
         self._type_name = sensor_type_name
         self._type = sensor_type
-        self._name = 'IGD {}'.format(sensor_type['name'])
+        self._name = '{} {}'.format(device.name, sensor_type['name'])
         self._state = None
 
     @property
@@ -88,6 +86,11 @@ class RawIGDSensor(Entity):
         """Return the name of the sensor."""
         return self._name
 
+    @property
+    def unique_id(self) -> str:
+        """Return an unique ID."""
+        return '{}_{}'.format(self._device.udn, self._type_name)
+
     @property
     def state(self) -> str:
         """Return the state of the device."""
@@ -116,143 +119,121 @@ class RawIGDSensor(Entity):
             self._state = await self._device.async_get_total_packets_sent()
 
 
-class KBytePerSecondIGDSensor(Entity):
+class PerSecondIGDSensor(Entity):
+    """Abstract representation of a X Sent/Received per second sensor."""
+
+    def __init__(self, device, direction):
+        """Initializer."""
+        self._device = device
+        self._direction = direction
+
+        self._state = None
+        self._last_value = None
+        self._last_update_time = None
+
+    @property
+    def unit(self) -> str:
+        """Unit we are measuring in."""
+        raise NotImplementedError()
+
+    @property
+    def _async_fetch_value(self):
+        """Fetch a value from the IGD."""
+        raise NotImplementedError()
+
+    @property
+    def unique_id(self) -> str:
+        """Return an unique ID."""
+        return '{}_{}/sec_{}'.format(self._device.udn,
+                                     self.unit,
+                                     self._direction)
+
+    @property
+    def name(self) -> str:
+        """Return the name of the sensor."""
+        return '{} {}/sec {}'.format(self._device.name,
+                                     self.unit,
+                                     self._direction)
+
+    @property
+    def icon(self) -> str:
+        """Icon to use in the frontend, if any."""
+        return 'mdi:server-network'
+
+    @property
+    def unit_of_measurement(self) -> str:
+        """Return the unit of measurement of this entity, if any."""
+        return '{}/sec'.format(self.unit)
+
+    def _is_overflowed(self, new_value) -> bool:
+        """Check if value has overflowed."""
+        return new_value < self._last_value
+
+    async def async_update(self):
+        """Get the latest information from the IGD."""
+        _LOGGER.debug('%s: async_update', self)
+        new_value = await self._async_fetch_value()
+
+        if self._last_value is None:
+            self._last_value = new_value
+            self._last_update_time = datetime.now()
+            return
+
+        now = datetime.now()
+        if self._is_overflowed(new_value):
+            self._state = None  # temporarily report nothing
+        else:
+            delta_time = (now - self._last_update_time).seconds
+            delta_value = new_value - self._last_value
+            self._state = (delta_value / delta_time)
+
+        self._last_value = new_value
+        self._last_update_time = now
+
+
+class KBytePerSecondIGDSensor(PerSecondIGDSensor):
     """Representation of a KBytes Sent/Received per second sensor."""
 
-    def __init__(self, device, direction):
-        """Initializer."""
-        self._device = device
-        self._direction = direction
-
-        self._last_value = None
-        self._last_update_time = None
-        self._state = None
-
     @property
-    def unique_id(self) -> str:
+    def unit(self) -> str:
-        """Return an unique ID."""
+        """Unit we are measuring in."""
-        return '{}:kbytes_{}'.format(self._device.udn, self._direction)
+        return 'kbyte'
 
-    @property
+    async def _async_fetch_value(self) -> float:
-    def name(self) -> str:
+        """"""
-        """Return the name of the sensor."""
+        if self._direction == IN:
-        return '{} kbytes/sec {}'.format(self._device.name,
+            return await self._device.async_get_total_bytes_received()
-                                             self._direction)
+
+        return await self._device.async_get_total_bytes_sent()
 
     @property
     def state(self):
         """Return the state of the device."""
-        return self._state
+        if self._state is None:
+            return None
 
-    @property
+        return format(float(self._state / KBYTE), '.1f')
-    def icon(self) -> str:
-        """Icon to use in the frontend, if any."""
-        return 'mdi:server-network'
-
-    @property
-    def unit_of_measurement(self) -> str:
-        """Return the unit of measurement of this entity, if any."""
-        return 'kbytes/sec'
-
-    def _is_overflowed(self, new_value) -> bool:
-        """Check if value has overflowed."""
-        return new_value < self._last_value
-
-    async def async_update(self):
-        """Get the latest information from the IGD."""
-        _LOGGER.debug('%s: async_update', self)
-
-        if self._direction == IN:
-            new_value = await self._device.async_get_total_bytes_received()
-        else:
-            new_value = await self._device.async_get_total_bytes_sent()
-
-        if self._last_value is None:
-            self._last_value = new_value
-            self._last_update_time = datetime.now()
-            return
-
-        now = datetime.now()
-        if self._is_overflowed(new_value):
-            _LOGGER.debug('%s: Overflow: old value: %s, new value: %s',
-                          self, self._last_value, new_value)
-            self._state = None  # temporarily report nothing
-        else:
-            delta_time = (now - self._last_update_time).seconds
-            delta_value = new_value - self._last_value
-            value = (delta_value / delta_time) / KBYTE
-            self._state = format(float(value), '.1f')
-
-        self._last_value = new_value
-        self._last_update_time = now
 
 
-class PacketsPerSecondIGDSensor(Entity):
+class PacketsPerSecondIGDSensor(PerSecondIGDSensor):
     """Representation of a Packets Sent/Received per second sensor."""
 
-    def __init__(self, device, direction):
-        """Initializer."""
-        self._device = device
-        self._direction = direction
-
-        self._last_value = None
-        self._last_update_time = None
-        self._state = None
-
     @property
-    def unique_id(self) -> str:
+    def unit(self) -> str:
-        """Return an unique ID."""
+        """Unit we are measuring in."""
-        return '{}:packets_{}'.format(self._device.udn, self._direction)
+        return 'packets'
 
-    @property
+    async def _async_fetch_value(self) -> float:
-    def name(self) -> str:
+        """"""
-        """Return the name of the sensor."""
+        if self._direction == IN:
-        return '{} packets/sec {}'.format(self._device.name,
+            return await self._device.async_get_total_packets_received()
-                                              self._direction)
+
+        return await self._device.async_get_total_packets_sent()
 
     @property
     def state(self):
         """Return the state of the device."""
-        return self._state
+        if self._state is None:
+            return None
 
-    @property
+        return format(float(self._state), '.1f')
-    def icon(self) -> str:
-        """Icon to use in the frontend, if any."""
-        return 'mdi:server-network'
-
-    @property
-    def unit_of_measurement(self) -> str:
-        """Return the unit of measurement of this entity, if any."""
-        return '#/sec'
-
-    def _is_overflowed(self, new_value) -> bool:
-        """Check if value has overflowed."""
-        return new_value < self._last_value
-
-    async def async_update(self):
-        """Get the latest information from the IGD."""
-        _LOGGER.debug('%s: async_update', self)
-
-        if self._direction == IN:
-            new_value = await self._device.async_get_total_bytes_received()
-        else:
-            new_value = await self._device.async_get_total_bytes_sent()
-
-        if self._last_value is None:
-            self._last_value = new_value
-            self._last_update_time = datetime.now()
-            return
-
-        now = datetime.now()
-        if self._is_overflowed(new_value):
-            _LOGGER.debug('%s: Overflow: old value: %s, new value: %s',
-                          self, self._last_value, new_value)
-            self._state = None  # temporarily report nothing
-        else:
-            delta_time = (now - self._last_update_time).seconds
-            delta_value = new_value - self._last_value
-            value = delta_value / delta_time
-            self._state = format(float(value), '.1f')
-
-        self._last_value = new_value
-        self._last_update_time = now

tests/components/igd/test_config_flow.py

@@ -10,6 +10,9 @@ async def test_flow_none_discovered(hass):
     """Test no device discovered flow."""
     flow = igd_config_flow.IgdFlowHandler()
     flow.hass = hass
+    hass.data[igd.DOMAIN] = {
+        'discovered': {}
+    }
 
     result = await flow.async_step_user()
     assert result['type'] == 'abort'

tests/components/igd/test_init.py

@@ -93,7 +93,7 @@ async def test_async_setup_entry_default(hass):
         hass.bus.async_fire(EVENT_HOMEASSISTANT_STOP)
         await hass.async_block_till_done()
 
-    assert hass.data[igd.DOMAIN]['devices'][udn] is None
+    assert udn not in hass.data[igd.DOMAIN]['devices']
     assert len(mock_igd_device.async_add_port_mapping.mock_calls) == 0
     assert len(mock_igd_device.async_delete_port_mapping.mock_calls) == 0
 
@@ -128,6 +128,6 @@ async def test_async_setup_entry_port_forward(hass):
         hass.bus.async_fire(EVENT_HOMEASSISTANT_STOP)
         await hass.async_block_till_done()
 
-    assert hass.data[igd.DOMAIN]['devices'][udn] is None
+    assert udn not in hass.data[igd.DOMAIN]['devices']
     assert len(mock_igd_device.async_add_port_mapping.mock_calls) > 0
     assert len(mock_igd_device.async_delete_port_mapping.mock_calls) > 0