core/homeassistant/components/sensor/speedtest.py

"""
homeassistant.components.sensor.speedtest
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Speedtest.net sensor based on speedtest-cli.

For more details about this platform, please refer to the documentation at
https://home-assistant.io/components/sensor.speedtest/
"""
import logging
import sys
import re
from datetime import timedelta
from subprocess import check_output
from homeassistant.util import Throttle
from homeassistant.helpers.entity import Entity
from homeassistant.helpers.event import track_time_change
from homeassistant.components.sensor import DOMAIN
import homeassistant.util.dt as dt_util

REQUIREMENTS = ['speedtest-cli==0.3.4']
_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]+')

CONF_MONITORED_CONDITIONS = 'monitored_conditions'
CONF_MINUTE = 'minute'
CONF_HOUR = 'hour'
CONF_DAY = 'day'
SENSOR_TYPES = {
    'ping': ['Ping', 'ms'],
    'download': ['Download', 'Mbit/s'],
    'upload': ['Upload', 'Mbit/s'],
}

# Return cached results if last scan was less then this time ago
MIN_TIME_BETWEEN_UPDATES = timedelta(minutes=1)


def setup_platform(hass, config, add_devices, discovery_info=None):
    """ Setup the Speedtest sensor. """

    data = SpeedtestData(hass, config)
    dev = []
    for sensor in config[CONF_MONITORED_CONDITIONS]:
        if sensor not in SENSOR_TYPES:
            _LOGGER.error('Sensor type: "%s" does not exist', sensor)
        else:
            dev.append(SpeedtestSensor(data, sensor))

    add_devices(dev)

    def update(call=None):
        """ Update service for manual updates. """
        data.update(dt_util.now())
        for sensor in dev:
            sensor.update()

    hass.services.register(DOMAIN, 'update_speedtest', update)


# pylint: disable=too-few-public-methods
class SpeedtestSensor(Entity):
    """ Implements a speedtest.net sensor. """

    def __init__(self, speedtest_data, sensor_type):
        self._name = SENSOR_TYPES[sensor_type][0]
        self.speedtest_client = speedtest_data
        self.type = sensor_type
        self._state = None
        self._unit_of_measurement = SENSOR_TYPES[self.type][1]

    @property
    def name(self):
        return '{} {}'.format('Speedtest', self._name)

    @property
    def state(self):
        """ Returns the state of the device. """
        return self._state

    @property
    def unit_of_measurement(self):
        """ Unit of measurement of this entity, if any. """
        return self._unit_of_measurement

    def update(self):
        """ Gets the latest data from Forecast.io and updates the states. """
        data = self.speedtest_client.data
        if data is not None:
            if self.type == 'ping':
                self._state = data['ping']
            elif self.type == 'download':
                self._state = data['download']
            elif self.type == 'upload':
                self._state = data['upload']


class SpeedtestData(object):
    """ Gets the latest data from speedtest.net. """

    def __init__(self, hass, config):
        self.data = None
        self.hass = hass
        self.path = hass.config.path
        track_time_change(self.hass, self.update,
                          minute=config.get(CONF_MINUTE, 0),
                          hour=config.get(CONF_HOUR, None),
                          day=config.get(CONF_DAY, None))

    @Throttle(MIN_TIME_BETWEEN_UPDATES)
    def update(self, now):
        """ Gets the latest data from speedtest.net. """
        _LOGGER.info('Executing speedtest')
        re_output = _SPEEDTEST_REGEX.split(
            check_output([sys.executable, self.path(
                'lib', 'speedtest_cli.py'), '--simple']).decode("utf-8"))
        self.data = {'ping': round(float(re_output[1]), 2),
                     'download': round(float(re_output[2]), 2),
                     'upload': round(float(re_output[3]), 2)}