From 01602d5439a1fc0724acec8003f7eecb74865fa0 Mon Sep 17 00:00:00 2001
From: MrS0m30n3 <ytubedlg@gmail.com>
Date: Thu, 24 Nov 2016 15:01:52 +0200
Subject: [PATCH] DownloadItem: Add 'Error' stage

---
 tests/test_ditem.py               | 51 ++++++++++++++++++++++++++++---
 youtube_dl_gui/downloadmanager.py | 13 ++++++--
 2 files changed, 58 insertions(+), 6 deletions(-)

diff --git a/tests/test_ditem.py b/tests/test_ditem.py
index 6c8176b..ac8635d 100644
--- a/tests/test_ditem.py
+++ b/tests/test_ditem.py
@@ -118,6 +118,10 @@ class TestSetItemStage(unittest.TestCase):
         self.assertEqual(self.ditem.stage, "Paused")
         self.assertEqual(self.ditem.progress_stats["status"], "Paused")
 
+        self.ditem.stage = "Error"
+        self.assertEqual(self.ditem.stage, "Error")
+        self.assertEqual(self.ditem.progress_stats["status"], "Error")
+
     def test_set_stage_invalid(self):
         raised = False
 
@@ -228,18 +232,21 @@ class TestDownloadItemPrivate(unittest.TestCase):
         ditem = DownloadItem("url", ["-f", "flv"])
 
         active_status = ["Pre Processing", "Downloading", "Post Processing"]
-        complete_status = ["Finished", "Error", "Warning", "Stopped", "Already Downloaded", "Filesize Abort"]
+        complete_status = ["Finished", "Warning", "Already Downloaded"]
+        error_status = ["Error", "Stopped", "Filesize Abort"]
 
         for status in active_status:
-            ditem._stage = "Queued"
             ditem._set_stage(status)
             self.assertEqual(ditem.stage, "Active")
 
         for status in complete_status:
-            ditem._stage = "Active"
             ditem._set_stage(status)
             self.assertEqual(ditem.stage, "Completed")
 
+        for status in error_status:
+            ditem._set_stage(status)
+            self.assertEqual(ditem.stage, "Error")
+
 
 class TestReset(unittest.TestCase):
 
@@ -248,9 +255,45 @@ class TestReset(unittest.TestCase):
     def setUp(self):
         self.ditem = DownloadItem("url", ["-f", "flv"])
 
-    def test_reset_error_status(self):
+    def test_reset_completed_stage(self):
         self.ditem._stage = "Completed"
         self.ditem.path = os.path.join("/home", "user")
+        self.ditem.filenames = ["file"]
+        self.ditem.extensions = [".mp4"]
+        self.ditem.progress_stats = {
+            "filename": "file",
+            "extension": ".mp4",
+            "filsize": "6.66MiB",
+            "percent": "100%",
+            "speed": "-",
+            "eta": "00:00",
+            "status": "Finished",
+            "playlist_size": "",
+            "playlist_index": ""
+        }
+
+        self.ditem.reset()
+
+        self.assertEqual(self.ditem._stage, "Queued")
+        self.assertEqual(self.ditem.path, "")
+        self.assertEqual(self.ditem.filenames, [])
+        self.assertEqual(self.ditem.extensions, [])
+        self.assertEqual(
+            self.ditem.progress_stats,
+            {"filename": "url",
+             "extension": "-",
+             "filesize": "-",
+             "percent": "0%",
+             "speed": "-",
+             "eta": "-",
+             "status": "Queued",
+             "playlist_size": "",
+             "playlist_index": ""}
+        )
+
+    def test_reset_error_stage(self):
+        self.ditem._stage = "Error"
+        self.ditem.path = os.path.join("/home", "user")
         self.ditem.filenames = ["file1", "file2", "file"]
         self.ditem.extensions = [".mp4", ".m4a", ".mp4"]
         self.ditem.progress_stats = {
diff --git a/youtube_dl_gui/downloadmanager.py b/youtube_dl_gui/downloadmanager.py
index 09e2de2..5c50216 100644
--- a/youtube_dl_gui/downloadmanager.py
+++ b/youtube_dl_gui/downloadmanager.py
@@ -73,6 +73,8 @@ class DownloadItem(object):
 
         COMPLETED_STAGES (tuple): Sub stages of the 'Completed' stage.
 
+        ERROR_STAGES (tuple): Sub stages of the 'Error' stage.
+
     Args:
         url (string): URL that corresponds to the download item.
 
@@ -80,11 +82,13 @@ class DownloadItem(object):
 
     """
 
-    STAGES = ("Queued", "Active", "Paused", "Completed")
+    STAGES = ("Queued", "Active", "Paused", "Completed", "Error")
 
     ACTIVE_STAGES = ("Pre Processing", "Downloading", "Post Processing")
 
-    COMPLETED_STAGES = ("Finished", "Error", "Warning", "Stopped", "Already Downloaded", "Filesize Abort")
+    COMPLETED_STAGES = ("Finished", "Warning", "Already Downloaded")
+
+    ERROR_STAGES = ("Error", "Stopped", "Filesize Abort")
 
     def __init__(self, url, options):
         self.url = url
@@ -110,6 +114,8 @@ class DownloadItem(object):
             self.progress_stats["status"] = self.COMPLETED_STAGES[0]
         if value == "Paused":
             self.progress_stats["status"] = value
+        if value == "Error":
+            self.progress_stats["status"] = self.ERROR_STAGES[0]
 
         self._stage = value
 
@@ -179,6 +185,9 @@ class DownloadItem(object):
         if status in self.COMPLETED_STAGES:
             self._stage = self.STAGES[3]
 
+        if status in self.ERROR_STAGES:
+            self._stage = self.STAGES[4]
+
     def __eq__(self, other):
         return self.object_id == other.object_id