From fc45670686179d4f703ec0dff6b7f125c47476bb Mon Sep 17 00:00:00 2001
From: Pascal Vizeli <pascal.vizeli@syshack.ch>
Date: Thu, 7 Mar 2019 21:09:43 +0100
Subject: [PATCH] Fix bug with update (#938)

---
 hassio/addons/addon.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/hassio/addons/addon.py b/hassio/addons/addon.py
index 85165e76f..3963b65eb 100644
--- a/hassio/addons/addon.py
+++ b/hassio/addons/addon.py
@@ -122,7 +122,10 @@ class Addon(CoreSysAttributes):
     def _set_update(self, image: str, version: str) -> None:
         """Update version of add-on."""
         self._data.system[self._id] = deepcopy(self._data.cache[self._id])
-        self._data.user[self._id][ATTR_VERSION] = version
+        self._data.user[self._id].update({
+            ATTR_VERSION: version,
+            ATTR_IMAGE: image,
+        })
         self.save_data()
 
     def _restore_data(self, user: Dict[str, Any], system: Dict[str, Any], image: str) -> None: