Merge pull request #402 from hzxjy1/zipTest

Close zipfile hander manually and add a test

nhentai/downloader.py

@@ -139,6 +139,7 @@ class Downloader(Singleton):
             except EnvironmentError as e:
                 logger.critical(str(e))
         self.folder:str = folder
+        self.close = lambda: None # Only available in class CompressedDownloader
 
     def start_download(self, queue, folder='') -> bool:
         if not isinstance(folder, (str,)):
@@ -164,6 +165,8 @@ class Downloader(Singleton):
         # Prevent coroutines infection
         asyncio.run(self.fiber(coroutines))
 
+        self.close()
+
         return True
 
 class CompressedDownloader(Downloader):
@@ -171,6 +174,7 @@ class CompressedDownloader(Downloader):
         filename = f'{folder}.zip'
         print(filename)
         self.zipfile = zipfile.ZipFile(filename,'w')
+        self.close = lambda: self.zipfile.close()
 
     async def save(self, filename, response) -> bool:
         if response is None:
@@ -189,6 +193,3 @@ class CompressedDownloader(Downloader):
         image_data.seek(0)
         self.zipfile.writestr(filename, image_data.read())
         return True
-
-    def __del__(self):
-         self.zipfile.close()

tests/test_download.py

@@ -1,14 +1,27 @@
 import unittest
 import os
+import zipfile
 import urllib3.exceptions
 
 from nhentai import constant
 from nhentai.cmdline import load_config
-from nhentai.downloader import Downloader
+from nhentai.downloader import Downloader, CompressedDownloader
 from nhentai.parser import doujinshi_parser
 from nhentai.doujinshi import Doujinshi
 from nhentai.utils import generate_html
 
+did = 440546
+
+def has_jepg_file(path):
+    with zipfile.ZipFile(path, 'r') as zf:
+        return '01.jpg' in zf.namelist()
+
+def is_zip_file(path):
+    try:
+        with zipfile.ZipFile(path, 'r') as _:
+            return True
+    except (zipfile.BadZipFile, FileNotFoundError):
+        return False
 
 class TestDownload(unittest.TestCase):
     def setUp(self) -> None:
@@ -17,17 +30,27 @@ class TestDownload(unittest.TestCase):
         constant.CONFIG['cookie'] = os.getenv('NHENTAI_COOKIE')
         constant.CONFIG['useragent'] = os.getenv('NHENTAI_UA')
 
+        self.info = Doujinshi(**doujinshi_parser(did), name_format='%i')
+
     def test_download(self):
-        did = 440546
+        info = self.info
-        info = Doujinshi(**doujinshi_parser(did), name_format='%i')
         info.downloader = Downloader(path='/tmp', threads=5)
         info.download()
 
-        self.assertTrue(os.path.exists(f'/tmp/{did}/001.jpg'))
+        self.assertTrue(os.path.exists(f'/tmp/{did}/01.jpg'))
 
         generate_html('/tmp', info)
         self.assertTrue(os.path.exists(f'/tmp/{did}/index.html'))
 
+    def test_zipfile_download(self):
+        info = self.info
+        info.downloader = CompressedDownloader(path='/tmp', threads=5)
+        info.download()
+
+        zipfile_path = f'/tmp/{did}.zip'
+        self.assertTrue(os.path.exists(zipfile_path))
+        self.assertTrue(is_zip_file(zipfile_path))
+        self.assertTrue(has_jepg_file(zipfile_path))
 
 if __name__ == '__main__':
     unittest.main()