AirVisual: Entity Registry updates and cleanup (#12319)

* AirVisual: Entity Registry updates and cleanup

* Small cleanup

* Owner-requested changes

* Changed hashing function

* Put a better class instatiation mechanism in place

* Small cleanup

* Reverting unintended breaking change

* Removing hashing as part of creating the unique ID

* Attempting to jumpstart Travis

homeassistant/components/sensor/airvisual.py

@@ -19,7 +19,6 @@ from homeassistant.helpers.entity import Entity
 from homeassistant.util import Throttle
 
 REQUIREMENTS = ['pyairvisual==1.0.0']
-
 _LOGGER = getLogger(__name__)
 
 ATTR_CITY = 'city'
@@ -27,7 +26,6 @@ ATTR_COUNTRY = 'country'
 ATTR_POLLUTANT_SYMBOL = 'pollutant_symbol'
 ATTR_POLLUTANT_UNIT = 'pollutant_unit'
 ATTR_REGION = 'region'
-ATTR_TIMESTAMP = 'timestamp'
 
 CONF_CITY = 'city'
 CONF_COUNTRY = 'country'
@@ -39,6 +37,12 @@ VOLUME_MICROGRAMS_PER_CUBIC_METER = 'µg/m3'
 
 MIN_TIME_BETWEEN_UPDATES = timedelta(minutes=10)
 
+SENSOR_TYPES = [
+    ('AirPollutionLevelSensor', 'Air Pollution Level', 'mdi:scale'),
+    ('AirQualityIndexSensor', 'Air Quality Index', 'mdi:format-list-numbers'),
+    ('MainPollutantSensor', 'Main Pollutant', 'mdi:chemical-weapon'),
+]
+
 POLLUTANT_LEVEL_MAPPING = [
     {'label': 'Good', 'minimum': 0, 'maximum': 50},
     {'label': 'Moderate', 'minimum': 51, 'maximum': 100},
@@ -58,11 +62,6 @@ POLLUTANT_MAPPING = {
 }
 
 SENSOR_LOCALES = {'cn': 'Chinese', 'us': 'U.S.'}
-SENSOR_TYPES = [
-    ('AirPollutionLevelSensor', 'Air Pollution Level', 'mdi:scale'),
-    ('AirQualityIndexSensor', 'Air Quality Index', 'mdi:format-list-numbers'),
-    ('MainPollutantSensor', 'Main Pollutant', 'mdi:chemical-weapon'),
-]
 
 PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend({
     vol.Required(CONF_API_KEY): cv.string,
@@ -80,7 +79,13 @@ PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend({
 
 def setup_platform(hass, config, add_devices, discovery_info=None):
     """Configure the platform and add the sensors."""
-    import pyairvisual as pav
+    from pyairvisual import Client
+
+    classes = {
+        'AirPollutionLevelSensor': AirPollutionLevelSensor,
+        'AirQualityIndexSensor': AirQualityIndexSensor,
+        'MainPollutantSensor': MainPollutantSensor
+    }
 
     api_key = config.get(CONF_API_KEY)
     monitored_locales = config.get(CONF_MONITORED_CONDITIONS)
@@ -95,60 +100,63 @@ def setup_platform(hass, config, add_devices, discovery_info=None):
     if city and state and country:
         _LOGGER.debug(
             "Using city, state, and country: %s, %s, %s", city, state, country)
+        location_id = ','.join((city, state, country))
         data = AirVisualData(
-            pav.Client(api_key), city=city, state=state, country=country,
+            Client(api_key), city=city, state=state, country=country,
             show_on_map=show_on_map)
     else:
         _LOGGER.debug(
             "Using latitude and longitude: %s, %s", latitude, longitude)
+        location_id = ','.join((str(latitude), str(longitude)))
         data = AirVisualData(
-            pav.Client(api_key), latitude=latitude, longitude=longitude,
+            Client(api_key), latitude=latitude, longitude=longitude,
             radius=radius, show_on_map=show_on_map)
 
     data.update()
+
     sensors = []
     for locale in monitored_locales:
         for sensor_class, name, icon in SENSOR_TYPES:
-            sensors.append(globals()[sensor_class](data, name, icon, locale))
+            sensors.append(classes[sensor_class](
+                data,
+                name,
+                icon,
+                locale,
+                location_id
+            ))
 
     add_devices(sensors, True)
 
 
-def merge_two_dicts(dict1, dict2):
-    """Merge two dicts into a new dict as a shallow copy."""
-    final = dict1.copy()
-    final.update(dict2)
-    return final
-
-
 class AirVisualBaseSensor(Entity):
     """Define a base class for all of our sensors."""
 
-    def __init__(self, data, name, icon, locale):
+    def __init__(self, data, name, icon, locale, entity_id):
         """Initialize the sensor."""
         self.data = data
+        self._attrs = {}
         self._icon = icon
         self._locale = locale
         self._name = name
         self._state = None
+        self._entity_id = entity_id
         self._unit = None
 
     @property
     def device_state_attributes(self):
         """Return the device state attributes."""
-        attrs = merge_two_dicts({
+        self._attrs.update({
             ATTR_ATTRIBUTION: CONF_ATTRIBUTION,
-            ATTR_TIMESTAMP: self.data.pollution_info.get('ts')
+        })
-        }, self.data.attrs)
 
         if self.data.show_on_map:
-            attrs[ATTR_LATITUDE] = self.data.latitude
+            self._attrs[ATTR_LATITUDE] = self.data.latitude
-            attrs[ATTR_LONGITUDE] = self.data.longitude
+            self._attrs[ATTR_LONGITUDE] = self.data.longitude
         else:
-            attrs['lati'] = self.data.latitude
+            self._attrs['lati'] = self.data.latitude
-            attrs['long'] = self.data.longitude
+            self._attrs['long'] = self.data.longitude
 
-        return attrs
+        return self._attrs
 
     @property
     def icon(self):
@@ -169,6 +177,11 @@ class AirVisualBaseSensor(Entity):
 class AirPollutionLevelSensor(AirVisualBaseSensor):
     """Define a sensor to measure air pollution level."""
 
+    @property
+    def unique_id(self):
+        """Return a unique, HASS-friendly identifier for this entity."""
+        return '{0}_pollution_level'.format(self._entity_id)
+
     def update(self):
         """Update the status of the sensor."""
         self.data.update()
@@ -189,6 +202,11 @@ class AirPollutionLevelSensor(AirVisualBaseSensor):
 class AirQualityIndexSensor(AirVisualBaseSensor):
     """Define a sensor to measure AQI."""
 
+    @property
+    def unique_id(self):
+        """Return a unique, HASS-friendly identifier for this entity."""
+        return '{0}_aqi'.format(self._entity_id)
+
     @property
     def unit_of_measurement(self):
         """Return the unit the value is expressed in."""
@@ -205,19 +223,16 @@ class AirQualityIndexSensor(AirVisualBaseSensor):
 class MainPollutantSensor(AirVisualBaseSensor):
     """Define a sensor to the main pollutant of an area."""
 
-    def __init__(self, data, name, icon, locale):
+    def __init__(self, data, name, icon, locale, entity_id):
         """Initialize the sensor."""
-        super().__init__(data, name, icon, locale)
+        super().__init__(data, name, icon, locale, entity_id)
         self._symbol = None
         self._unit = None
 
     @property
-    def device_state_attributes(self):
+    def unique_id(self):
-        """Return the device state attributes."""
+        """Return a unique, HASS-friendly identifier for this entity."""
-        return merge_two_dicts(super().device_state_attributes, {
+        return '{0}_main_pollutant'.format(self._entity_id)
-            ATTR_POLLUTANT_SYMBOL: self._symbol,
-            ATTR_POLLUTANT_UNIT: self._unit
-        })
 
     def update(self):
         """Update the status of the sensor."""
@@ -229,6 +244,11 @@ class MainPollutantSensor(AirVisualBaseSensor):
         self._unit = pollution_info.get('unit')
         self._symbol = symbol
 
+        self._attrs.update({
+            ATTR_POLLUTANT_SYMBOL: self._symbol,
+            ATTR_POLLUTANT_UNIT: self._unit
+        })
+
 
 class AirVisualData(object):
     """Define an object to hold sensor data."""
@@ -252,7 +272,7 @@ class AirVisualData(object):
     @Throttle(MIN_TIME_BETWEEN_UPDATES)
     def update(self):
         """Update with new AirVisual data."""
-        import pyairvisual.exceptions as exceptions
+        from pyairvisual.exceptions import HTTPError
 
         try:
             if self.city and self.state and self.country:
@@ -272,7 +292,7 @@ class AirVisualData(object):
                 ATTR_REGION: resp.get('state'),
                 ATTR_COUNTRY: resp.get('country')
             }
-        except exceptions.HTTPError as exc_info:
+        except HTTPError as exc_info:
             _LOGGER.error("Unable to retrieve data on this location: %s",
                           self.__dict__)
             _LOGGER.debug(exc_info)