fix pickle (closes #814)

app/tests/test_persistent_queue.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 import os
 import tempfile
 import unittest
+from unittest.mock import patch
 
 from ytdl import DownloadInfo, PersistentQueue
 
@@ -71,6 +72,48 @@ class PersistentQueueTests(unittest.TestCase):
             pq2.load()
             self.assertTrue(pq2.exists("http://load.example"))
 
+    def test_put_rollbacks_in_memory_queue_when_shelf_write_fails(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            path = os.path.join(tmp, "queue")
+            pq = PersistentQueue("queue", path)
+            dl = _FakeDownload(_make_info("http://rollback.example"))
+            self.assertFalse(pq.exists("http://rollback.example"))
+
+            orig_open = __import__("shelve").open
+
+            def bad_open(filename, flag="c", *args, **kwargs):
+                if flag == "w":
+                    raise OSError("simulated shelf failure")
+                return orig_open(filename, flag, *args, **kwargs)
+
+            with patch("ytdl.shelve.open", bad_open):
+                with self.assertRaises(OSError):
+                    pq.put(dl)
+
+            self.assertFalse(pq.exists("http://rollback.example"))
+
+    def test_put_rollbacks_to_previous_download_when_replace_fails(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            path = os.path.join(tmp, "queue")
+            pq = PersistentQueue("queue", path)
+            first = _FakeDownload(_make_info("http://same.example"))
+            second = _FakeDownload(_make_info("http://same.example"))
+            second.info.title = "Replaced title"
+            pq.put(first)
+
+            orig_open = __import__("shelve").open
+
+            def bad_open(filename, flag="c", *args, **kwargs):
+                if flag == "w":
+                    raise OSError("simulated shelf failure")
+                return orig_open(filename, flag, *args, **kwargs)
+
+            with patch("ytdl.shelve.open", bad_open):
+                with self.assertRaises(OSError):
+                    pq.put(second)
+
+            self.assertEqual(pq.get("http://same.example").info.title, "Title")
+
 
 if __name__ == "__main__":
     unittest.main()

app/tests/test_ytdl_utils.py

@@ -2,15 +2,17 @@
 
 from __future__ import annotations
 
+import pickle
 import tempfile
+import threading
 import unittest
 from pathlib import Path
 
 from ytdl import (
     DownloadInfo,
-    _convert_generators_to_lists,
     _convert_srt_to_txt_file,
     _outtmpl_substitute_field,
+    _sanitize_entry_for_pickle,
     _sanitize_path_component,
 )
 
@@ -35,17 +37,41 @@ class OuttmplSubstituteFieldTests(unittest.TestCase):
         self.assertEqual(_outtmpl_substitute_field("%(other)s", "title", "x"), "%(other)s")
 
 
-class ConvertGeneratorsToListsTests(unittest.TestCase):
+class SanitizeEntryForPickleTests(unittest.TestCase):
     def test_nested(self):
         def g():
             yield 1
 
         obj = {"a": g(), "b": [g()]}
-        out = _convert_generators_to_lists(obj)
+        out = _sanitize_entry_for_pickle(obj)
         self.assertEqual(out, {"a": [1], "b": [[1]]})
+        pickle.dumps(out)
 
     def test_plain(self):
-        self.assertEqual(_convert_generators_to_lists(5), 5)
+        self.assertEqual(_sanitize_entry_for_pickle(5), 5)
+
+    def test_set_converted_to_list(self):
+        obj = {"s": {1, 2}}
+        out = _sanitize_entry_for_pickle(obj)
+        self.assertEqual(sorted(out["s"]), [1, 2])
+        pickle.dumps(out)
+
+    def test_map_iterator(self):
+        out = _sanitize_entry_for_pickle({"m": map(int, ["1", "2"])})
+        self.assertEqual(out, {"m": [1, 2]})
+
+    def test_lock_replaced_with_none(self):
+        lock = threading.Lock()
+        out = _sanitize_entry_for_pickle({"k": lock})
+        self.assertIsNone(out["k"])
+        pickle.dumps(out)
+
+    def test_ordered_dict(self):
+        from collections import OrderedDict
+
+        od = OrderedDict([("z", 1), ("a", 2)])
+        out = _sanitize_entry_for_pickle(od)
+        self.assertEqual(out, {"z": 1, "a": 2})
 
 
 class ConvertSrtToTxtTests(unittest.TestCase):