Refactor the restore backup implementation to RestoreBackupJob

3 files changed, 106 insertions, 63 deletions

diff --git a/plugins/CuraDrive/src/CreateBackupJob.py b/plugins/CuraDrive/src/CreateBackupJob.py
index 603733137b..4a5764a3b6 100644
--- a/plugins/CuraDrive/src/CreateBackupJob.py
+++ b/plugins/CuraDrive/src/CreateBackupJob.py
@@ -36,10 +36,8 @@ class CreateBackupJob(Job):
         self._api_backup_url = api_backup_url
         self._jsonCloudScope = JsonDecoratorScope(UltimakerCloudScope(CuraApplication.getInstance()))
 
-
         self._backup_zip = None
-        self._upload_success = False
-        self._upload_success_available = threading.Event()
+        self._job_done = threading.Event()
         self.backup_upload_error_message = ""
 
     def run(self) -> None:
@@ -62,7 +60,7 @@ class CreateBackupJob(Job):
         backup_meta_data["description"] = "{}.backup.{}.cura.zip".format(timestamp, backup_meta_data["cura_release"])
         self._requestUploadSlot(backup_meta_data, len(self._backup_zip))
 
-        self._upload_success_available.wait()
+        self._job_done.wait()
         upload_message.hide()
 
     def _requestUploadSlot(self, backup_metadata: Dict[str, Any], backup_size: int) -> None:
@@ -89,12 +87,12 @@ class CreateBackupJob(Job):
         if error is not None:
             Logger.warning(str(error))
             self.backup_upload_error_message = "Could not upload backup."
-            self._upload_success_available.set()
+            self._job_done.set()
             return
         if reply.attribute(QNetworkRequest.HttpStatusCodeAttribute) >= 300:
             Logger.warning("Could not request backup upload: %s", HttpRequestManager.readText(reply))
             self.backup_upload_error_message = "Could not upload backup."
-            self._upload_success_available.set()
+            self._job_done.set()
             return
 
         backup_upload_url = HttpRequestManager.readJSON(reply)["data"]["upload_url"]
@@ -119,4 +117,4 @@ class CreateBackupJob(Job):
             Message(catalog.i18nc("@info:backup_status", "There was an error while uploading your backup."),
                     title=self.MESSAGE_TITLE).show()
 
-        self._upload_success_available.set()
+        self._job_done.set()
diff --git a/plugins/CuraDrive/src/DriveApiService.py b/plugins/CuraDrive/src/DriveApiService.py
index 04f935268b..3af461126a 100644
--- a/plugins/CuraDrive/src/DriveApiService.py
+++ b/plugins/CuraDrive/src/DriveApiService.py
@@ -11,6 +11,7 @@ from UM.Signal import Signal, signalemitter
 from UM.TaskManagement.HttpRequestManager import HttpRequestManager
 from UM.TaskManagement.HttpRequestScope import JsonDecoratorScope
 from cura.CuraApplication import CuraApplication
+from plugins.CuraDrive.src.RestoreBackupJob import RestoreBackupJob
 from plugins.Toolbox.src.UltimakerCloudScope import UltimakerCloudScope
 
 from PyQt5.QtNetwork import QNetworkReply, QNetworkRequest
@@ -27,7 +28,6 @@ class DriveApiService:
     """The DriveApiService is responsible for interacting with the CuraDrive API and Cura's backup handling."""
 
     BACKUP_URL = "{}/backups".format(Settings.DRIVE_API_URL)
-    DISK_WRITE_BUFFER_SIZE = 512 * 1024
 
     restoringStateChanged = Signal()
     """Emits signal when restoring backup started or finished."""
@@ -82,61 +82,16 @@ class DriveApiService:
             # If there is no download URL, we can't restore the backup.
             return self._emitRestoreError()
 
-        def finishedCallback(reply: QNetworkReply, bu=backup) -> None:
-            self._onRestoreRequestCompleted(reply, None, bu)
+        restore_backup_job = RestoreBackupJob(backup)
+        restore_backup_job.finished.connect(self._onRestoreFinished)
+        restore_backup_job.start()
 
-        HttpRequestManager.getInstance().get(
-            url = download_url,
-            callback = finishedCallback,
-            error_callback = self._onRestoreRequestCompleted
-        )
-
-    def _onRestoreRequestCompleted(self, reply: QNetworkReply, error: Optional["QNetworkReply.NetworkError"] = None, backup = None):
-        if not HttpRequestManager.replyIndicatesSuccess(reply, error):
-            Logger.log("w",
-                       "Requesting backup failed, response code %s while trying to connect to %s",
-                       reply.attribute(QNetworkRequest.HttpStatusCodeAttribute), reply.url())
-            self._emitRestoreError()
-            return
-
-        # We store the file in a temporary path fist to ensure integrity.
-        temporary_backup_file = NamedTemporaryFile(delete = False)
-        with open(temporary_backup_file.name, "wb") as write_backup:
-            app = CuraApplication.getInstance()
-            bytes_read = reply.read(DriveApiService.DISK_WRITE_BUFFER_SIZE)
-            while bytes_read:
-                write_backup.write(bytes_read)
-                bytes_read = reply.read(DriveApiService.DISK_WRITE_BUFFER_SIZE)
-                app.processEvents()
-
-        if not self._verifyMd5Hash(temporary_backup_file.name, backup.get("md5_hash", "")):
-            # Don't restore the backup if the MD5 hashes do not match.
-            # This can happen if the download was interrupted.
-            Logger.log("w", "Remote and local MD5 hashes do not match, not restoring backup.")
-            return self._emitRestoreError()
-
-        # Tell Cura to place the backup back in the user data folder.
-        with open(temporary_backup_file.name, "rb") as read_backup:
-            self._cura_api.backups.restoreBackup(read_backup.read(), backup.get("metadata", {}))
-            self.restoringStateChanged.emit(is_restoring = False)
-
-    def _emitRestoreError(self) -> None:
-        self.restoringStateChanged.emit(is_restoring = False,
-                                        error_message = catalog.i18nc("@info:backup_status",
-                                                                         "There was an error trying to restore your backup."))
-
-    @staticmethod
-    def _verifyMd5Hash(file_path: str, known_hash: str) -> bool:
-        """Verify the MD5 hash of a file.
-
-        :param file_path: Full path to the file.
-        :param known_hash: The known MD5 hash of the file.
-        :return: Success or not.
-        """
-
-        with open(file_path, "rb") as read_backup:
-            local_md5_hash = base64.b64encode(hashlib.md5(read_backup.read()).digest(), altchars = b"_-").decode("utf-8")
-            return known_hash == local_md5_hash
+    def _onRestoreFinished(self, job: "RestoreBackupJob"):
+        if job.restore_backup_error_message != "":
+            # If the job contains an error message we pass it along so the UI can display it.
+            self.restoringStateChanged.emit(is_restoring=False)
+        else:
+            self.restoringStateChanged.emit(is_restoring = False, error_message = job.restore_backup_error_message)
 
     def deleteBackup(self, backup_id: str, finishedCallable: Callable[[bool], None]):
 
@@ -153,5 +108,6 @@ class DriveApiService:
             scope= self._jsonCloudScope
         )
 
-    def _onDeleteRequestCompleted(self, reply: QNetworkReply, error: Optional["QNetworkReply.NetworkError"] = None, callable = None):
+    @staticmethod
+    def _onDeleteRequestCompleted(reply: QNetworkReply, error: Optional["QNetworkReply.NetworkError"] = None, callable = None):
         callable(HttpRequestManager.replyIndicatesSuccess(reply, error))
diff --git a/plugins/CuraDrive/src/RestoreBackupJob.py b/plugins/CuraDrive/src/RestoreBackupJob.py
new file mode 100644
index 0000000000..86e32aca97
--- /dev/null
+++ b/plugins/CuraDrive/src/RestoreBackupJob.py
@@ -0,0 +1,89 @@
+import base64
+import hashlib
+import threading
+from tempfile import NamedTemporaryFile
+from typing import Optional, Any, Dict
+
+from PyQt5.QtNetwork import QNetworkReply, QNetworkRequest
+
+from UM.Job import Job
+from UM.Logger import Logger
+from UM.PackageManager import catalog
+from UM.TaskManagement.HttpRequestManager import HttpRequestManager
+from cura.CuraApplication import CuraApplication
+
+
+class RestoreBackupJob(Job):
+    """Downloads a backup and overwrites local configuration with the backup.
+
+     When `Job.finished` emits, `restore_backup_error_message` will either be `""` (no error) or an error message
+     """
+
+    DISK_WRITE_BUFFER_SIZE = 512 * 1024
+    DEFAULT_ERROR_MESSAGE = catalog.i18nc("@info:backup_status", "There was an error trying to restore your backup.")
+
+    def __init__(self, backup: Dict[str, Any]) -> None:
+        """ Create a new restore Job. start the job by calling start()
+
+        :param backup: A dict containing a backup spec
+        """
+
+        super().__init__()
+        self._job_done = threading.Event()
+
+        self._backup = backup
+        self.restore_backup_error_message = ""
+
+    def run(self):
+
+        HttpRequestManager.getInstance().get(
+            url = self._backup.get("download_url"),
+            callback = self._onRestoreRequestCompleted,
+            error_callback = self._onRestoreRequestCompleted
+        )
+
+        self._job_done.wait()
+
+    def _onRestoreRequestCompleted(self, reply: QNetworkReply, error: Optional["QNetworkReply.NetworkError"] = None):
+        if not HttpRequestManager.replyIndicatesSuccess(reply, error):
+            Logger.warning("Requesting backup failed, response code %s while trying to connect to %s",
+                           reply.attribute(QNetworkRequest.HttpStatusCodeAttribute), reply.url())
+            self.restore_backup_error_message = self.DEFAULT_ERROR_MESSAGE
+            self._job_done.set()
+            return
+
+        # We store the file in a temporary path fist to ensure integrity.
+        temporary_backup_file = NamedTemporaryFile(delete = False)
+        with open(temporary_backup_file.name, "wb") as write_backup:
+            app = CuraApplication.getInstance()
+            bytes_read = reply.read(self.DISK_WRITE_BUFFER_SIZE)
+            while bytes_read:
+                write_backup.write(bytes_read)
+                bytes_read = reply.read(self.DISK_WRITE_BUFFER_SIZE)
+                app.processEvents()
+
+        if not self._verifyMd5Hash(temporary_backup_file.name, self._backup.get("md5_hash", "")):
+            # Don't restore the backup if the MD5 hashes do not match.
+            # This can happen if the download was interrupted.
+            Logger.log("w", "Remote and local MD5 hashes do not match, not restoring backup.")
+            self.restore_backup_error_message = self.DEFAULT_ERROR_MESSAGE
+
+        # Tell Cura to place the backup back in the user data folder.
+        with open(temporary_backup_file.name, "rb") as read_backup:
+            cura_api = CuraApplication.getInstance().getCuraAPI()
+            cura_api.backups.restoreBackup(read_backup.read(), self._backup.get("metadata", {}))
+
+        self._job_done.set()
+
+    @staticmethod
+    def _verifyMd5Hash(file_path: str, known_hash: str) -> bool:
+        """Verify the MD5 hash of a file.
+
+        :param file_path: Full path to the file.
+        :param known_hash: The known MD5 hash of the file.
+        :return: Success or not.
+        """
+
+        with open(file_path, "rb") as read_backup:
+            local_md5_hash = base64.b64encode(hashlib.md5(read_backup.read()).digest(), altchars = b"_-").decode("utf-8")
+            return known_hash == local_md5_hash