From 3d1a324f33493ba6477e6caf61248a0a6160b7bc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Joakim=20S=C3=B8rensen?= <joasoe@gmail.com>
Date: Thu, 1 Nov 2018 09:20:30 +0100
Subject: [PATCH] Add functionality to the version sensor (#18067)

* Added functionality to the version sensor.

* Corrected typo.

* Change default name to not cause a breaking change.

* Use vol.lower in the schema.

* Add missing blank line.

* Change order of cv.string and vol.Lower.
---
 homeassistant/components/sensor/version.py | 78 +++++++++++++++++++---
 requirements_all.txt                       |  3 +
 2 files changed, 70 insertions(+), 11 deletions(-)

diff --git a/homeassistant/components/sensor/version.py b/homeassistant/components/sensor/version.py
index eba4b1b8350..b71ae158181 100644
--- a/homeassistant/components/sensor/version.py
+++ b/homeassistant/components/sensor/version.py
@@ -1,54 +1,110 @@
 """
-Support for displaying the current version of Home Assistant.
+Sensor that can display the current Home Assistant versions.
 
 For more details about this platform, please refer to the documentation at
 https://home-assistant.io/components/sensor.version/
 """
 import logging
+from datetime import timedelta
 
 import voluptuous as vol
 
 import homeassistant.helpers.config_validation as cv
+from homeassistant.helpers.aiohttp_client import async_get_clientsession
 from homeassistant.components.sensor import PLATFORM_SCHEMA
-from homeassistant.const import __version__, CONF_NAME
+from homeassistant.const import CONF_NAME, CONF_SOURCE
 from homeassistant.helpers.entity import Entity
+from homeassistant.util import Throttle
+
+REQUIREMENTS = ['pyhaversion==2.0.1']
 
 _LOGGER = logging.getLogger(__name__)
 
+CONF_BETA = 'beta'
+CONF_IMAGE = 'image'
+
+DEFAULT_IMAGE = 'default'
 DEFAULT_NAME = "Current Version"
+DEFAULT_SOURCE = 'local'
+
+ICON = 'mdi:package-up'
+
+TIME_BETWEEN_UPDATES = timedelta(minutes=5)
 
 PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend({
+    vol.Optional(CONF_BETA, default=False): cv.boolean,
+    vol.Optional(CONF_IMAGE, default=DEFAULT_IMAGE): vol.All(cv.string,
+                                                             vol.Lower),
     vol.Optional(CONF_NAME, default=DEFAULT_NAME): cv.string,
+    vol.Optional(CONF_SOURCE, default=DEFAULT_SOURCE): vol.All(cv.string,
+                                                               vol.Lower),
 })
 
 
 async def async_setup_platform(
         hass, config, async_add_entities, discovery_info=None):
     """Set up the Version sensor platform."""
+    from pyhaversion import Version
+    beta = config.get(CONF_BETA)
+    image = config.get(CONF_IMAGE)
     name = config.get(CONF_NAME)
+    source = config.get(CONF_SOURCE)
 
-    async_add_entities([VersionSensor(name)])
+    session = async_get_clientsession(hass)
+    if beta:
+        branch = 'beta'
+    else:
+        branch = 'stable'
+    haversion = VersionData(Version(hass.loop, session, branch, image), source)
+
+    async_add_entities([VersionSensor(haversion, name)], True)
 
 
 class VersionSensor(Entity):
     """Representation of a Home Assistant version sensor."""
 
-    def __init__(self, name):
+    def __init__(self, haversion, name):
         """Initialize the Version sensor."""
+        self.haversion = haversion
         self._name = name
-        self._state = __version__
+        self._state = None
+
+    async def async_update(self):
+        """Get the latest version information."""
+        await self.haversion.async_update()
 
     @property
     def name(self):
         """Return the name of the sensor."""
         return self._name
 
-    @property
-    def should_poll(self):
-        """No polling needed."""
-        return False
-
     @property
     def state(self):
         """Return the state of the sensor."""
-        return self._state
+        return self.haversion.api.version
+
+    @property
+    def device_state_attributes(self):
+        """Return attributes for the sensor."""
+        return self.haversion.api.version_data
+
+
+class VersionData:
+    """Get the latest data and update the states."""
+
+    def __init__(self, api, source):
+        """Initialize the data object."""
+        self.api = api
+        self.source = source
+
+    @Throttle(TIME_BETWEEN_UPDATES)
+    async def async_update(self):
+        """Get the latest version information."""
+        if self.source == 'pypi':
+            await self.api.get_pypi_version()
+        elif self.source == 'hassio':
+            await self.api.get_hassio_version()
+        elif self.source == 'docker':
+            await self.api.get_docker_version()
+        else:
+            await self.api.get_local_version()
diff --git a/requirements_all.txt b/requirements_all.txt
index 45305c26e3b..d6c8965d31a 100644
--- a/requirements_all.txt
+++ b/requirements_all.txt
@@ -919,6 +919,9 @@ pygtfs-homeassistant==0.1.3.dev0
 # homeassistant.components.remote.harmony
 pyharmony==1.0.20
 
+# homeassistant.components.sensor.version
+pyhaversion==2.0.1
+
 # homeassistant.components.binary_sensor.hikvision
 pyhik==0.1.8