Add attributes (fixes #12332) (#12377)

* Add attributes (fixes #12332)

* Fix pylint issue

homeassistant/components/sensor/speedtest.py

@@ -6,27 +6,30 @@ https://home-assistant.io/components/sensor.speedtest/
 """
 import asyncio
 import logging
-import re
-import sys
-from subprocess import check_output, CalledProcessError
 
 import voluptuous as vol
 
-import homeassistant.util.dt as dt_util
+from homeassistant.components.sensor import DOMAIN, PLATFORM_SCHEMA
+from homeassistant.const import ATTR_ATTRIBUTION, CONF_MONITORED_CONDITIONS
 import homeassistant.helpers.config_validation as cv
-from homeassistant.components.sensor import (DOMAIN, PLATFORM_SCHEMA)
-from homeassistant.const import CONF_MONITORED_CONDITIONS
 from homeassistant.helpers.entity import Entity
 from homeassistant.helpers.event import track_time_change
 from homeassistant.helpers.restore_state import async_get_last_state
+import homeassistant.util.dt as dt_util
 
-REQUIREMENTS = ['speedtest-cli==1.0.7']
+REQUIREMENTS = ['speedtest-cli==2.0.0']
 
 _LOGGER = logging.getLogger(__name__)
-_SPEEDTEST_REGEX = re.compile(r'Ping:\s(\d+\.\d+)\sms[\r\n]+'
-                              r'Download:\s(\d+\.\d+)\sMbit/s[\r\n]+'
-                              r'Upload:\s(\d+\.\d+)\sMbit/s[\r\n]+')
 
+ATTR_BYTES_RECEIVED = 'bytes_received'
+ATTR_BYTES_SENT = 'bytes_sent'
+ATTR_SERVER_COUNTRY = 'server_country'
+ATTR_SERVER_HOST = 'server_host'
+ATTR_SERVER_ID = 'server_id'
+ATTR_SERVER_LATENCY = 'latency'
+ATTR_SERVER_NAME = 'server_name'
+
+CONF_ATTRIBUTION = "Data retrieved from Speedtest by Ookla"
 CONF_SECOND = 'second'
 CONF_MINUTE = 'minute'
 CONF_HOUR = 'hour'
@@ -45,28 +48,26 @@ SENSOR_TYPES = {
 PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend({
     vol.Required(CONF_MONITORED_CONDITIONS):
         vol.All(cv.ensure_list, [vol.In(list(SENSOR_TYPES))]),
-    vol.Optional(CONF_SERVER_ID): cv.positive_int,
-    vol.Optional(CONF_SECOND, default=[0]):
-        vol.All(cv.ensure_list, [vol.All(vol.Coerce(int), vol.Range(0, 59))]),
-    vol.Optional(CONF_MINUTE, default=[0]):
-        vol.All(cv.ensure_list, [vol.All(vol.Coerce(int), vol.Range(0, 59))]),
-    vol.Optional(CONF_HOUR):
-        vol.All(cv.ensure_list, [vol.All(vol.Coerce(int), vol.Range(0, 23))]),
     vol.Optional(CONF_DAY):
         vol.All(cv.ensure_list, [vol.All(vol.Coerce(int), vol.Range(1, 31))]),
+    vol.Optional(CONF_HOUR):
+        vol.All(cv.ensure_list, [vol.All(vol.Coerce(int), vol.Range(0, 23))]),
     vol.Optional(CONF_MANUAL, default=False): cv.boolean,
+    vol.Optional(CONF_MINUTE, default=[0]):
+        vol.All(cv.ensure_list, [vol.All(vol.Coerce(int), vol.Range(0, 59))]),
+    vol.Optional(CONF_SECOND, default=[0]):
+        vol.All(cv.ensure_list, [vol.All(vol.Coerce(int), vol.Range(0, 59))]),
+    vol.Optional(CONF_SERVER_ID): cv.positive_int,
 })
 
 
 def setup_platform(hass, config, add_devices, discovery_info=None):
     """Set up the Speedtest sensor."""
     data = SpeedtestData(hass, config)
+
     dev = []
     for sensor in config[CONF_MONITORED_CONDITIONS]:
-        if sensor not in SENSOR_TYPES:
+        dev.append(SpeedtestSensor(data, sensor))
-            _LOGGER.error("Sensor type: %s does not exist", sensor)
-        else:
-            dev.append(SpeedtestSensor(data, sensor))
 
     add_devices(dev)
 
@@ -88,6 +89,7 @@ class SpeedtestSensor(Entity):
         self.speedtest_client = speedtest_data
         self.type = sensor_type
         self._state = None
+        self._data = None
         self._unit_of_measurement = SENSOR_TYPES[self.type][1]
 
     @property
@@ -110,18 +112,32 @@ class SpeedtestSensor(Entity):
         """Return icon."""
         return ICON
 
+    @property
+    def device_state_attributes(self):
+        """Return the state attributes."""
+        if self._data is not None:
+            return {
+                ATTR_ATTRIBUTION: CONF_ATTRIBUTION,
+                ATTR_BYTES_RECEIVED: self._data['bytes_received'],
+                ATTR_BYTES_SENT: self._data['bytes_sent'],
+                ATTR_SERVER_COUNTRY: self._data['server']['country'],
+                ATTR_SERVER_ID: self._data['server']['id'],
+                ATTR_SERVER_LATENCY: self._data['server']['latency'],
+                ATTR_SERVER_NAME: self._data['server']['name'],
+            }
+
     def update(self):
         """Get the latest data and update the states."""
-        data = self.speedtest_client.data
+        self._data = self.speedtest_client.data
-        if data is None:
+        if self._data is None:
             return
 
         if self.type == 'ping':
-            self._state = data['ping']
+            self._state = self._data['ping']
         elif self.type == 'download':
-            self._state = data['download']
+            self._state = round(self._data['download'] / 10**6, 2)
         elif self.type == 'upload':
-            self._state = data['upload']
+            self._state = round(self._data['upload'] / 10**6, 2)
 
     @asyncio.coroutine
     def async_added_to_hass(self):
@@ -148,20 +164,14 @@ class SpeedtestData(object):
     def update(self, now):
         """Get the latest data from speedtest.net."""
         import speedtest
+        _LOGGER.debug("Executing speedtest...")
 
-        _LOGGER.info("Executing speedtest...")
+        servers = [] if self._server_id is None else [self._server_id]
-        try:
-            args = [sys.executable, speedtest.__file__, '--simple']
-            if self._server_id:
-                args = args + ['--server', str(self._server_id)]
 
-            re_output = _SPEEDTEST_REGEX.split(
+        speed = speedtest.Speedtest()
-                check_output(args).decode('utf-8'))
+        speed.get_servers(servers)
-        except CalledProcessError as process_error:
+        speed.get_best_server()
-            _LOGGER.error("Error executing speedtest: %s", process_error)
+        speed.download()
-            return
+        speed.upload()
-        self.data = {
+
-            'ping': round(float(re_output[1]), 2),
+        self.data = speed.results.dict()
-            'download': round(float(re_output[2]), 2),
-            'upload': round(float(re_output[3]), 2),
-        }

requirements_all.txt

@@ -1121,7 +1121,7 @@ snapcast==2.0.8
 somecomfort==0.5.0
 
 # homeassistant.components.sensor.speedtest
-speedtest-cli==1.0.7
+speedtest-cli==2.0.0
 
 # homeassistant.components.recorder
 # homeassistant.scripts.db_migrator