From e5fc6846e0ab8e8a152d47d053f7cdd847631e20 Mon Sep 17 00:00:00 2001
From: pvizeli <pvizeli@syshack.ch>
Date: Tue, 11 Jul 2017 16:52:16 +0200
Subject: [PATCH] fix snapshot with custom image

---
 hassio/homeassistant.py      |  3 +++
 hassio/snapshots/__init__.py |  7 +++----
 hassio/snapshots/snapshot.py | 31 ++++++++++++++++++++++++++++++-
 hassio/snapshots/validate.py |  3 ++-
 4 files changed, 38 insertions(+), 6 deletions(-)

diff --git a/hassio/homeassistant.py b/hassio/homeassistant.py
index e7c6721e3..c062a6be8 100644
--- a/hassio/homeassistant.py
+++ b/hassio/homeassistant.py
@@ -73,9 +73,12 @@ class HomeAssistant(JsonConfig):
         if image is None and version is None:
             self._data.pop(ATTR_IMAGE, None)
             self._data.pop(ATTR_VERSION, None)
+
+            self.docker.image = self.image
         else:
             if image:
                 self._data[ATTR_IMAGE] = image
+                self.docker.image = image
             if version:
                 self._data[ATTR_VERSION] = version
         self.save()
diff --git a/hassio/snapshots/__init__.py b/hassio/snapshots/__init__.py
index 8d2bd3f91..bc8603a9b 100644
--- a/hassio/snapshots/__init__.py
+++ b/hassio/snapshots/__init__.py
@@ -46,8 +46,7 @@ class SnapshotsManager(object):
         snapshot.create(slug, name, date_str, sys_type)
 
         # set general data
-        snapshot.homeassistant_version = self.homeassistant.version
-        snapshot.homeassistant_devices = self.homeassistant.devices
+        snapshot.snapshot_homeassistant(self.homeassistant)
         snapshot.repositories = self.config.addons_repositories
 
         return snapshot
@@ -198,7 +197,7 @@ class SnapshotsManager(object):
                 await snapshot.restore_folders()
 
                 # start homeassistant restore
-                self.homeassistant.devices = snapshot.homeassistant_devices
+                snapshot.restore_homeassistant(self.homeassistant)
                 task_hass = self.loop.create_task(
                     self.homeassistant.update(snapshot.homeassistant_version))
 
@@ -280,7 +279,7 @@ class SnapshotsManager(object):
                     await snapshot.restore_folders(folders)
 
                 if homeassistant:
-                    self.homeassistant.devices = snapshot.homeassistant_devices
+                    snapshot.restore_homeassistant(self.homeassistant)
                     tasks.append(self.homeassistant.update(
                         snapshot.homeassistant_version))
 
diff --git a/hassio/snapshots/snapshot.py b/hassio/snapshots/snapshot.py
index 1882a5f07..a8730e20a 100644
--- a/hassio/snapshots/snapshot.py
+++ b/hassio/snapshots/snapshot.py
@@ -13,7 +13,8 @@ from .validate import SCHEMA_SNAPSHOT, ALL_FOLDERS
 from .util import remove_folder
 from ..const import (
     ATTR_SLUG, ATTR_NAME, ATTR_DATE, ATTR_ADDONS, ATTR_REPOSITORIES,
-    ATTR_HOMEASSISTANT, ATTR_FOLDERS, ATTR_VERSION, ATTR_TYPE, ATTR_DEVICES)
+    ATTR_HOMEASSISTANT, ATTR_FOLDERS, ATTR_VERSION, ATTR_TYPE, ATTR_DEVICES,
+    ATTR_IMAGE)
 from ..tools import write_json_file
 
 _LOGGER = logging.getLogger(__name__)
@@ -90,6 +91,16 @@ class Snapshot(object):
         """Set snapshot homeassistant devices."""
         self._data[ATTR_HOMEASSISTANT][ATTR_DEVICES] = value
 
+    @property
+    def homeassistant_image(self):
+        """Return snapshot homeassistant custom image."""
+        return self._data[ATTR_HOMEASSISTANT].get(ATTR_IMAGE)
+
+    @homeassistant_image.setter
+    def homeassistant_image(self, value):
+        """Set snapshot homeassistant custom image."""
+        self._data[ATTR_HOMEASSISTANT][ATTR_IMAGE] = value
+
     @property
     def size(self):
         """Return snapshot size."""
@@ -111,6 +122,24 @@ class Snapshot(object):
         self._data[ATTR_REPOSITORIES] = []
         self._data[ATTR_FOLDERS] = []
 
+    def snapshot_homeassistant(self, homeassistant):
+        """Read all data from homeassistant object."""
+        self.homeassistant_version = homeassistant.version
+        self.homeassistant_devices = homeassistant.devices
+
+        # custom image
+        if self.homeassistant.is_custom_image:
+            self.homeassistant_image = homeassistant.image
+
+    def restore_homeassistant(self, homeassistant):
+        """Write all data to homeassistant object."""
+        homeassistant.devices = self.homeassistant_devices
+
+        # custom image
+        if self.homeassistant_image:
+            homeassistant.set_custom(
+                self.homeassistant_image, self.homeassistant_version)
+
     async def load(self):
         """Read snapshot.json from tar file."""
         if not self.tar_file.is_file():
diff --git a/hassio/snapshots/validate.py b/hassio/snapshots/validate.py
index 2259009d5..e0a497d76 100644
--- a/hassio/snapshots/validate.py
+++ b/hassio/snapshots/validate.py
@@ -5,7 +5,7 @@ import voluptuous as vol
 from ..const import (
     ATTR_REPOSITORIES, ATTR_ADDONS, ATTR_NAME, ATTR_SLUG, ATTR_DATE,
     ATTR_VERSION, ATTR_HOMEASSISTANT, ATTR_FOLDERS, ATTR_TYPE, ATTR_DEVICES,
-    FOLDER_SHARE, FOLDER_HOMEASSISTANT, FOLDER_ADDONS, FOLDER_SSL,
+    ATTR_IMAGE, FOLDER_SHARE, FOLDER_HOMEASSISTANT, FOLDER_ADDONS, FOLDER_SSL,
     SNAPSHOT_FULL, SNAPSHOT_PARTIAL)
 from ..validate import HASS_DEVICES
 
@@ -20,6 +20,7 @@ SCHEMA_SNAPSHOT = vol.Schema({
     vol.Required(ATTR_HOMEASSISTANT): vol.Schema({
         vol.Required(ATTR_VERSION): vol.Coerce(str),
         vol.Optional(ATTR_DEVICES, default=[]): HASS_DEVICES,
+        vol.Optional(ATTR_IMAGE): vol.Coerce(str),
     }),
     vol.Optional(ATTR_FOLDERS, default=[]): [vol.In(ALL_FOLDERS)],
     vol.Optional(ATTR_ADDONS, default=[]): [vol.Schema({