Split out iperf3 into a component with a sensor platform

.coveragerc

@@ -204,6 +204,7 @@ omit =
     homeassistant/components/insteon/*
     homeassistant/components/ios/*
     homeassistant/components/iota/*
+    homeassistant/components/iperf3/*
     homeassistant/components/isy994/*
     homeassistant/components/joaoapps_join/*
     homeassistant/components/juicenet/*
@@ -473,7 +474,6 @@ omit =
     homeassistant/components/sensor/imap_email_content.py
     homeassistant/components/sensor/imap.py
     homeassistant/components/sensor/influxdb.py
-    homeassistant/components/sensor/iperf3.py
     homeassistant/components/sensor/irish_rail_transport.py
     homeassistant/components/sensor/kwb.py
     homeassistant/components/sensor/lacrosse.py

homeassistant/components/iperf3/__init__.py

@@ -0,0 +1,185 @@
+"""Support for Iperf3 network measurement tool."""
+import logging
+from datetime import timedelta
+
+import voluptuous as vol
+
+import homeassistant.helpers.config_validation as cv
+from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN
+from homeassistant.const import CONF_MONITORED_CONDITIONS, CONF_PORT, \
+    CONF_HOST, CONF_PROTOCOL, CONF_HOSTS, CONF_SCAN_INTERVAL
+from homeassistant.helpers.discovery import async_load_platform
+from homeassistant.helpers.dispatcher import dispatcher_send
+from homeassistant.helpers.event import async_track_time_interval
+
+REQUIREMENTS = ['iperf3==0.1.10']
+
+DOMAIN = 'iperf3'
+DATA_UPDATED = '{}_data_updated'.format(DOMAIN)
+
+_LOGGER = logging.getLogger(__name__)
+
+CONF_DURATION = 'duration'
+CONF_PARALLEL = 'parallel'
+CONF_MANUAL = 'manual'
+
+DEFAULT_DURATION = 10
+DEFAULT_PORT = 5201
+DEFAULT_PARALLEL = 1
+DEFAULT_PROTOCOL = 'tcp'
+DEFAULT_INTERVAL = timedelta(minutes=60)
+
+ATTR_DOWNLOAD = 'download'
+ATTR_UPLOAD = 'upload'
+ATTR_VERSION = 'Version'
+ATTR_HOST = 'host'
+
+UNIT_OF_MEASUREMENT = 'Mbit/s'
+
+SENSOR_TYPES = {
+    ATTR_DOWNLOAD: [ATTR_DOWNLOAD.capitalize(), UNIT_OF_MEASUREMENT],
+    ATTR_UPLOAD: [ATTR_UPLOAD.capitalize(), UNIT_OF_MEASUREMENT],
+}
+
+PROTOCOLS = ['tcp', 'udp']
+
+HOST_CONFIG_SCHEMA = vol.Schema({
+    vol.Required(CONF_HOST): cv.string,
+    vol.Optional(CONF_PORT, default=DEFAULT_PORT): cv.port,
+    vol.Optional(CONF_DURATION, default=DEFAULT_DURATION): vol.Range(5, 10),
+    vol.Optional(CONF_PARALLEL, default=DEFAULT_PARALLEL): vol.Range(1, 20),
+    vol.Optional(CONF_PROTOCOL, default=DEFAULT_PROTOCOL): vol.In(PROTOCOLS),
+})
+
+CONFIG_SCHEMA = vol.Schema({
+    DOMAIN: vol.Schema({
+        vol.Required(CONF_HOSTS): vol.All(
+            cv.ensure_list, [HOST_CONFIG_SCHEMA]
+        ),
+        vol.Optional(CONF_MONITORED_CONDITIONS, default=list(SENSOR_TYPES)):
+            vol.All(cv.ensure_list, [vol.In(list(SENSOR_TYPES))]),
+        vol.Optional(CONF_SCAN_INTERVAL, default=DEFAULT_INTERVAL): vol.All(
+            cv.time_period, cv.positive_timedelta
+        ),
+        vol.Optional(CONF_MANUAL, default=False): cv.boolean,
+    })
+}, extra=vol.ALLOW_EXTRA)
+
+SERVICE_SCHEMA = vol.Schema({
+    vol.Optional(ATTR_HOST, default=None): cv.string,
+})
+
+
+async def async_setup(hass, config):
+    """Set up the iperf3 component."""
+    import iperf3
+
+    hass.data[DOMAIN] = {}
+
+    conf = config[DOMAIN]
+    for host in conf[CONF_HOSTS]:
+        host_name = host[CONF_HOST]
+
+        client = iperf3.Client()
+        client.duration = host[CONF_DURATION]
+        client.server_hostname = host_name
+        client.port = host[CONF_PORT]
+        client.num_streams = host[CONF_PARALLEL]
+        client.protocol = host[CONF_PROTOCOL]
+        client.verbose = False
+
+        data = hass.data[DOMAIN][host_name] = Iperf3Data(hass, client)
+
+        if not conf[CONF_MANUAL]:
+            async_track_time_interval(
+                hass, data.update, conf[CONF_SCAN_INTERVAL]
+            )
+
+    def update(call):
+        """Service call to manually update the data."""
+        called_host = call.data[ATTR_HOST]
+        if called_host in hass.data[DOMAIN]:
+            hass.data[DOMAIN][called_host].update()
+        else:
+            for iperf3_host in hass.data[DOMAIN].values():
+                iperf3_host.update()
+
+    hass.services.async_register(
+        DOMAIN, 'speedtest', update, schema=SERVICE_SCHEMA
+    )
+
+    hass.async_create_task(
+        async_load_platform(
+            hass,
+            SENSOR_DOMAIN,
+            DOMAIN,
+            conf[CONF_MONITORED_CONDITIONS],
+            config
+        )
+    )
+
+    return True
+
+
+class Iperf3Data:
+    """Get the latest data from iperf3."""
+
+    def __init__(self, hass, client):
+        """Initialize the data object."""
+        self._hass = hass
+        self._client = client
+        self.data = {
+            ATTR_DOWNLOAD: None,
+            ATTR_UPLOAD: None,
+            ATTR_VERSION: None
+        }
+
+    @property
+    def protocol(self):
+        """Return the protocol used for this connection."""
+        return self._client.protocol
+
+    @property
+    def host(self):
+        """Return the host connected to."""
+        return self._client.server_hostname
+
+    @property
+    def port(self):
+        """Return the port on the host connected to."""
+        return self._client.port
+
+    def update(self, now=None):
+        """Get the latest data from iperf3."""
+        if self.protocol == 'udp':
+            # UDP only have 1 way attribute
+            result = self._run_test(ATTR_DOWNLOAD)
+            self.data[ATTR_DOWNLOAD] = self.data[ATTR_UPLOAD] = getattr(
+                result, 'Mbps', None)
+            self.data[ATTR_VERSION] = getattr(result, 'version', None)
+        else:
+            result = self._run_test(ATTR_DOWNLOAD)
+            self.data[ATTR_DOWNLOAD] = getattr(
+                result, 'received_Mbps', None)
+            self.data[ATTR_VERSION] = getattr(result, 'version', None)
+            self.data[ATTR_UPLOAD] = getattr(
+                self._run_test(ATTR_UPLOAD), 'sent_Mbps', None)
+
+        dispatcher_send(self._hass, DATA_UPDATED, self.host)
+
+    def _run_test(self, test_type):
+        """Run and return the iperf3 data."""
+        self._client.reverse = test_type == ATTR_DOWNLOAD
+        try:
+            result = self._client.run()
+        except (AttributeError, OSError, ValueError) as error:
+            _LOGGER.error("Iperf3 error: %s", error)
+            return None
+
+        if result is not None and \
+           hasattr(result, 'error') and \
+           result.error is not None:
+            _LOGGER.error("Iperf3 error: %s", result.error)
+            return None
+
+        return result

homeassistant/components/iperf3/sensor.py

@@ -0,0 +1,100 @@
+"""Support for Iperf3 sensors."""
+from homeassistant.components.iperf3 import (
+    DATA_UPDATED, DOMAIN as IPERF3_DOMAIN, SENSOR_TYPES, ATTR_VERSION)
+from homeassistant.const import ATTR_ATTRIBUTION
+from homeassistant.core import callback
+from homeassistant.helpers.dispatcher import async_dispatcher_connect
+from homeassistant.helpers.restore_state import RestoreEntity
+
+DEPENDENCIES = ['iperf3']
+
+ATTRIBUTION = 'Data retrieved using Iperf3'
+
+ICON = 'mdi:speedometer'
+
+ATTR_PROTOCOL = 'Protocol'
+ATTR_REMOTE_HOST = 'Remote Server'
+ATTR_REMOTE_PORT = 'Remote Port'
+
+
+async def async_setup_platform(
+        hass, config, async_add_entities, discovery_info):
+    """Set up the Iperf3 sensor."""
+    sensors = []
+    for iperf3_host in hass.data[IPERF3_DOMAIN].values():
+        sensors.extend(
+            [Iperf3Sensor(iperf3_host, sensor) for sensor in discovery_info]
+        )
+    async_add_entities(sensors, True)
+
+
+class Iperf3Sensor(RestoreEntity):
+    """A Iperf3 sensor implementation."""
+
+    def __init__(self, iperf3_data, sensor_type):
+        """Initialize the sensor."""
+        self._name = \
+            "{} {}".format(SENSOR_TYPES[sensor_type][0], iperf3_data.host)
+        self._state = None
+        self._sensor_type = sensor_type
+        self._unit_of_measurement = SENSOR_TYPES[sensor_type][1]
+        self._iperf3_data = iperf3_data
+
+    @property
+    def name(self):
+        """Return the name of the sensor."""
+        return self._name
+
+    @property
+    def state(self):
+        """Return the state of the device."""
+        return self._state
+
+    @property
+    def unit_of_measurement(self):
+        """Return the unit of measurement of this entity, if any."""
+        return self._unit_of_measurement
+
+    @property
+    def icon(self):
+        """Return icon."""
+        return ICON
+
+    @property
+    def device_state_attributes(self):
+        """Return the state attributes."""
+        return {
+            ATTR_ATTRIBUTION: ATTRIBUTION,
+            ATTR_PROTOCOL: self._iperf3_data.protocol,
+            ATTR_REMOTE_HOST: self._iperf3_data.host,
+            ATTR_REMOTE_PORT: self._iperf3_data.port,
+            ATTR_VERSION: self._iperf3_data.data[ATTR_VERSION]
+        }
+
+    @property
+    def should_poll(self):
+        """Return the polling requirement for this sensor."""
+        return False
+
+    async def async_added_to_hass(self):
+        """Handle entity which will be added."""
+        await super().async_added_to_hass()
+        state = await self.async_get_last_state()
+        if not state:
+            return
+        self._state = state.state
+
+        async_dispatcher_connect(
+            self.hass, DATA_UPDATED, self._schedule_immediate_update
+        )
+
+    def update(self):
+        """Get the latest data and update the states."""
+        data = self._iperf3_data.data.get(self._sensor_type)
+        if data is not None:
+            self._state = round(data, 2)
+
+    @callback
+    def _schedule_immediate_update(self, host):
+        if host == self._iperf3_data.host:
+            self.async_schedule_update_ha_state(True)

homeassistant/components/iperf3/services.yaml

@@ -0,0 +1,6 @@
+speedtest:
+  description: Immediately take a speedest with iperf3
+  fields:
+    host:
+      description: The host name of the iperf3 server (already configured) to run a test with.
+      example: 'iperf.he.net'