fix:: always download the last file #81 (#86)

tangyoha · Jun 18, 2023 · d0d576f · d0d576f
1 parent cd0eac0
commit d0d576f
Show file tree

Hide file tree

Showing 3 changed files with 54 additions and 41 deletions.
diff --git a/module/app.py b/module/app.py
@@ -568,17 +568,25 @@ def update_config(self, immediate: bool = True):
         # pylint: disable = R1733
         for key, value in self.chat_download_config.items():
             # pylint: disable = W0201
-            self.chat_download_config[key].ids_to_retry = (
-                list(set(value.ids_to_retry) - set(value.downloaded_ids))
-                + value.failed_ids
+            before_last_read_message_id = self.config["chat"][idx].get(
+                "last_read_message_id", 0
             )
 
+            unfinished_ids = set(value.ids_to_retry) | set(
+                range(before_last_read_message_id, value.last_read_message_id + 1)
+            )
+            unfinished_ids -= set(value.downloaded_ids)
+
+            self.chat_download_config[key].ids_to_retry = list(unfinished_ids)
+
             if idx >= len(self.app_data["chat"]):
                 self.app_data["chat"].append({})
 
-            self.config["chat"][idx][
-                "last_read_message_id"
-            ] = value.last_read_message_id
+            if len(value.downloaded_ids) > 0:
+                self.config["chat"][idx]["last_read_message_id"] = (
+                    value.last_read_message_id + 1
+                )
+
             self.app_data["chat"][idx]["chat_id"] = key
             self.app_data["chat"][idx]["ids_to_retry"] = value.ids_to_retry
             idx += 1

diff --git a/tests/module/test_app.py b/tests/module/test_app.py
@@ -27,20 +27,23 @@ def test_app(self):
         self.assertEqual(app.proxy, {})
         self.assertEqual(app.restart_program, False)
 
-        app.last_read_message_id = 3
         app.chat_download_config[123] = ChatDownloadConfig()
-        app.chat_download_config[123].failed_ids.append(1)
-        app.chat_download_config[123].ids_to_retry.append(2)
-        app.config["chat"] = [{"chat_id": 123, "last_read_message_id": 0}]
+        app.chat_download_config[123].last_read_message_id = 13
+        app.chat_download_config[123].failed_ids.append(6)
+        app.chat_download_config[123].ids_to_retry.append(7)
+        app.chat_download_config[123].downloaded_ids.append(8)
+        app.chat_download_config[123].downloaded_ids.append(10)
+        app.chat_download_config[123].downloaded_ids.append(13)
+        app.config["chat"] = [{"chat_id": 123, "last_read_message_id": 5}]
 
         app.update_config(False)
 
         self.assertEqual(
-            app.chat_download_config[123].last_read_message_id,
+            app.chat_download_config[123].last_read_message_id + 1,
             app.config["chat"][0]["last_read_message_id"],
         )
         self.assertEqual(
-            app.chat_download_config[123].ids_to_retry,
+            [5, 6, 7, 9, 11, 12],
             app.app_data["chat"][0]["ids_to_retry"],
         )
 

diff --git a/tests/test_media_downloader.py b/tests/test_media_downloader.py
@@ -215,6 +215,36 @@ async def new_fetch_message(client: pyrogram.Client, message: pyrogram.types.Mes
     return message
 
 
+async def get_chat_history(client, *args, **kwargs):
+    items = [
+        MockMessage(
+            id=1213,
+            media=True,
+            voice=MockVoice(
+                mime_type="audio/ogg",
+                date=datetime(2019, 7, 25, 14, 53, 50),
+            ),
+        ),
+        MockMessage(
+            id=1214,
+            media=False,
+            text="test message 1",
+        ),
+        MockMessage(
+            id=1215,
+            media=False,
+            text="test message 2",
+        ),
+        MockMessage(
+            id=1216,
+            media=False,
+            text="test message 3",
+        ),
+    ]
+    for item in items:
+        yield item
+
+
 class MockClient:
     def __init__(self, *args, **kwargs):
         pass
@@ -228,35 +258,6 @@ async def start(self):
     async def stop(self):
         pass
 
-    async def get_chat_history(self, *args, **kwargs):
-        items = [
-            MockMessage(
-                id=1213,
-                media=True,
-                voice=MockVoice(
-                    mime_type="audio/ogg",
-                    date=datetime(2019, 7, 25, 14, 53, 50),
-                ),
-            ),
-            MockMessage(
-                id=1214,
-                media=False,
-                text="test message 1",
-            ),
-            MockMessage(
-                id=1215,
-                media=False,
-                text="test message 2",
-            ),
-            MockMessage(
-                id=1216,
-                media=False,
-                text="test message 3",
-            ),
-        ]
-        for item in items:
-            yield item
-
     async def get_messages(self, *args, **kwargs):
         if kwargs["message_ids"] == 7:
             return MockMessage(
@@ -359,6 +360,7 @@ async def edit_message_text(self, *args, **kwargs):
 
 @mock.patch("media_downloader.get_extension", new=get_extension)
 @mock.patch("media_downloader.fetch_message", new=new_fetch_message)
+@mock.patch("media_downloader.get_chat_history_v2", new=get_chat_history)
 @mock.patch("media_downloader.RETRY_TIME_OUT", new=0)
 class MediaDownloaderTestCase(unittest.TestCase):
     @classmethod