From 28d6910e56a5f7c186928cc7f0f4e650e6308802 Mon Sep 17 00:00:00 2001
From: Marcelo Moreira de Mello <tchello.mello@gmail.com>
Date: Sat, 26 May 2018 07:43:31 -0400
Subject: [PATCH] Added UDP and parallel streams support to Iperf3 (#14629)

---
 homeassistant/components/sensor/iperf3.py | 29 ++++++++++++++++++-----
 1 file changed, 23 insertions(+), 6 deletions(-)

diff --git a/homeassistant/components/sensor/iperf3.py b/homeassistant/components/sensor/iperf3.py
index 1a209faf17fb2c..8e030390f501b1 100644
--- a/homeassistant/components/sensor/iperf3.py
+++ b/homeassistant/components/sensor/iperf3.py
@@ -12,7 +12,7 @@
 from homeassistant.components.sensor import DOMAIN, PLATFORM_SCHEMA
 from homeassistant.const import (
     ATTR_ATTRIBUTION, ATTR_ENTITY_ID, CONF_MONITORED_CONDITIONS,
-    CONF_HOST, CONF_PORT)
+    CONF_HOST, CONF_PORT, CONF_PROTOCOL)
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.entity import Entity
 
@@ -27,13 +27,16 @@
 
 CONF_ATTRIBUTION = 'Data retrieved using Iperf3'
 CONF_DURATION = 'duration'
+CONF_PARALLEL = 'parallel'
 
 DEFAULT_DURATION = 10
 DEFAULT_PORT = 5201
+DEFAULT_PARALLEL = 1
+DEFAULT_PROTOCOL = 'tcp'
 
 IPERF3_DATA = 'iperf3'
 
-SCAN_INTERVAL = timedelta(minutes=30)
+SCAN_INTERVAL = timedelta(minutes=60)
 
 SERVICE_NAME = 'iperf3_update'
 
@@ -50,6 +53,9 @@
     vol.Required(CONF_PORT, default=DEFAULT_PORT): cv.port,
     vol.Required(CONF_HOST): cv.string,
     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(['tcp', 'udp']),
 })
 
 
@@ -70,6 +76,8 @@ def setup_platform(hass, config, add_devices, discovery_info=None):
             Iperf3Sensor(config[CONF_HOST],
                          config[CONF_PORT],
                          config[CONF_DURATION],
+                         config[CONF_PARALLEL],
+                         config[CONF_PROTOCOL],
                          sensor))
 
     hass.data[IPERF3_DATA]['sensors'].extend(dev)
@@ -98,10 +106,12 @@ def _service_handler(service):
 class Iperf3Sensor(Entity):
     """A Iperf3 sensor implementation."""
 
-    def __init__(self, server, port, duration, sensor_type):
+    def __init__(self, server, port, duration, streams,
+                 protocol, sensor_type):
         """Initialize the sensor."""
         self._attrs = {
             ATTR_ATTRIBUTION: CONF_ATTRIBUTION,
+            ATTR_PROTOCOL: protocol,
         }
         self._name = \
             "{} {}".format(SENSOR_TYPES[sensor_type][0], server)
@@ -111,6 +121,8 @@ def __init__(self, server, port, duration, sensor_type):
         self._port = port
         self._server = server
         self._duration = duration
+        self._num_streams = streams
+        self._protocol = protocol
         self.result = None
 
     @property
@@ -133,7 +145,6 @@ def device_state_attributes(self):
         """Return the state attributes."""
         if self.result is not None:
             self._attrs[ATTR_ATTRIBUTION] = CONF_ATTRIBUTION
-            self._attrs[ATTR_PROTOCOL] = self.result.protocol
             self._attrs[ATTR_REMOTE_HOST] = self.result.remote_host
             self._attrs[ATTR_REMOTE_PORT] = self.result.remote_port
             self._attrs[ATTR_VERSION] = self.result.version
@@ -147,6 +158,8 @@ def update(self):
         client.server_hostname = self._server
         client.port = self._port
         client.verbose = False
+        client.num_streams = self._num_streams
+        client.protocol = self._protocol
 
         # when testing download bandwith, reverse must be True
         if self._sensor_type == 'download':
@@ -154,7 +167,7 @@ def update(self):
 
         try:
             self.result = client.run()
-        except (OSError, AttributeError) as error:
+        except (AttributeError, OSError, ValueError) as error:
             self.result = None
             _LOGGER.error("Iperf3 sensor error: %s", error)
             return
@@ -166,7 +179,11 @@ def update(self):
             self.result = None
             return
 
-        if self._sensor_type == 'download':
+        # UDP only have 1 way attribute
+        if self._protocol == 'udp':
+            self._state = round(self.result.Mbps, 2)
+
+        elif self._sensor_type == 'download':
             self._state = round(self.result.received_Mbps, 2)
 
         elif self._sensor_type == 'upload':