From b1ef5042f91b4f431e6967cd726c9d79b88353dc Mon Sep 17 00:00:00 2001
From: Paulus Schoutsen <paulus@paulusschoutsen.nl>
Date: Wed, 30 Nov 2016 13:03:09 -0800
Subject: [PATCH] Make updater more robust (#4625)

---
 homeassistant/components/updater.py | 29 +++++++++++++++++++++--------
 tests/components/test_updater.py    | 13 ++++++++++++-
 2 files changed, 33 insertions(+), 9 deletions(-)

diff --git a/homeassistant/components/updater.py b/homeassistant/components/updater.py
index a4203876348..c05aedbb888 100644
--- a/homeassistant/components/updater.py
+++ b/homeassistant/components/updater.py
@@ -41,6 +41,11 @@ CONFIG_SCHEMA = vol.Schema({DOMAIN: {
     vol.Optional(CONF_REPORTING, default=True): cv.boolean
 }}, extra=vol.ALLOW_EXTRA)
 
+RESPONSE_SCHEMA = vol.Schema({
+    vol.Required('version'): str,
+    vol.Required('release-notes'): cv.url,
+})
+
 
 def _create_uuid(hass, filename=UPDATER_UUID_FILE):
     """Create UUID and save it in a file."""
@@ -83,7 +88,12 @@ def setup(hass, config):
 
 def check_newest_version(hass, huuid):
     """Check if a new version is available and report if one is."""
-    newest, releasenotes = get_newest_version(huuid)
+    result = get_newest_version(huuid)
+
+    if result is None:
+        return
+
+    newest, releasenotes = result
 
     if newest is None or 'dev' in CURRENT_VERSION:
         return
@@ -129,21 +139,24 @@ def get_newest_version(huuid):
     if not huuid:
         info_object = {}
 
+    res = None
     try:
         req = requests.post(UPDATER_URL, json=info_object, timeout=5)
         res = req.json()
+        res = RESPONSE_SCHEMA(res)
+
         _LOGGER.info(("Submitted analytics to Home Assistant servers. "
                       "Information submitted includes %s"), info_object)
         return (res['version'], res['release-notes'])
     except requests.RequestException:
-        _LOGGER.exception("Could not contact Home Assistant Update to check"
-                          "for updates")
+        _LOGGER.error("Could not contact Home Assistant Update to check "
+                      "for updates")
         return None
+
     except ValueError:
-        _LOGGER.exception("Received invalid response from Home Assistant"
-                          "Update")
+        _LOGGER.error("Received invalid response from Home Assistant Update")
         return None
-    except KeyError:
-        _LOGGER.exception("Response from Home Assistant Update did not"
-                          "include version")
+
+    except vol.Invalid:
+        _LOGGER.error('Got unexpected response: %s', res)
         return None
diff --git a/tests/components/test_updater.py b/tests/components/test_updater.py
index 94a43c1f281..8ca136bd8d7 100644
--- a/tests/components/test_updater.py
+++ b/tests/components/test_updater.py
@@ -6,6 +6,7 @@ import os
 
 import requests
 import requests_mock
+import voluptuous as vol
 
 from homeassistant.bootstrap import setup_component
 from homeassistant.components import updater
@@ -86,7 +87,7 @@ class TestUpdater(unittest.TestCase):
         mock_get.side_effect = ValueError
         self.assertIsNone(updater.get_newest_version(uuid))
 
-        mock_get.side_effect = KeyError
+        mock_get.side_effect = vol.Invalid('Expected dictionary')
         self.assertIsNone(updater.get_newest_version(uuid))
 
     def test_uuid_function(self):
@@ -119,3 +120,13 @@ class TestUpdater(unittest.TestCase):
 
         assert len(history) == 1
         assert history[0].json() == {}
+
+    @patch('homeassistant.components.updater.get_newest_version')
+    def test_error_during_fetch_works(
+            self, mock_get_newest_version):
+        """Test if no entity is created if same version."""
+        mock_get_newest_version.return_value = None
+
+        updater.check_newest_version(self.hass, None)
+
+        self.assertIsNone(self.hass.states.get(updater.ENTITY_ID))