mirror of
https://github.com/mikf/gallery-dl.git
synced 2024-11-21 18:22:30 +01:00
422 lines
11 KiB
Python
422 lines
11 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
|
|
# Copyright 2021-2023 Mike Fährmann
|
|
#
|
|
# This program is free software; you can redistribute it and/or modify
|
|
# it under the terms of the GNU General Public License version 2 as
|
|
# published by the Free Software Foundation.
|
|
|
|
import os
|
|
import sys
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
import io
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
from gallery_dl import job, config, text # noqa E402
|
|
from gallery_dl.extractor.common import Extractor, Message # noqa E402
|
|
|
|
|
|
class TestJob(unittest.TestCase):
|
|
|
|
def tearDown(self):
|
|
config.clear()
|
|
|
|
def _capture_stdout(self, extr_or_job):
|
|
if isinstance(extr_or_job, Extractor):
|
|
jobinstance = self.jobclass(extr_or_job)
|
|
else:
|
|
jobinstance = extr_or_job
|
|
|
|
with io.StringIO() as buffer:
|
|
stdout = sys.stdout
|
|
sys.stdout = buffer
|
|
try:
|
|
jobinstance.run()
|
|
finally:
|
|
sys.stdout = stdout
|
|
|
|
return buffer.getvalue()
|
|
|
|
|
|
class TestDownloadJob(TestJob):
|
|
jobclass = job.DownloadJob
|
|
|
|
def test_extractor_filter(self):
|
|
extr = TestExtractor.from_url("test:")
|
|
tjob = self.jobclass(extr)
|
|
|
|
func = tjob._build_extractor_filter()
|
|
self.assertEqual(func(TestExtractor) , False)
|
|
self.assertEqual(func(TestExtractorParent), False)
|
|
self.assertEqual(func(TestExtractorAlt) , True)
|
|
|
|
config.set((), "blacklist", ":test_subcategory")
|
|
func = tjob._build_extractor_filter()
|
|
self.assertEqual(func(TestExtractor) , False)
|
|
self.assertEqual(func(TestExtractorParent), True)
|
|
self.assertEqual(func(TestExtractorAlt) , False)
|
|
|
|
config.set((), "whitelist", "test_category:test_subcategory")
|
|
func = tjob._build_extractor_filter()
|
|
self.assertEqual(func(TestExtractor) , True)
|
|
self.assertEqual(func(TestExtractorParent), False)
|
|
self.assertEqual(func(TestExtractorAlt) , False)
|
|
|
|
|
|
class TestKeywordJob(TestJob):
|
|
jobclass = job.KeywordJob
|
|
|
|
def test_default(self):
|
|
self.maxDiff = None
|
|
extr = TestExtractor.from_url("test:self")
|
|
self.assertEqual(self._capture_stdout(extr), """\
|
|
Keywords for directory names:
|
|
-----------------------------
|
|
author['id']
|
|
123
|
|
author['name']
|
|
test
|
|
author['self']
|
|
<circular reference>
|
|
category
|
|
test_category
|
|
subcategory
|
|
test_subcategory
|
|
user['id']
|
|
123
|
|
user['name']
|
|
test
|
|
user['self']
|
|
<circular reference>
|
|
|
|
Keywords for filenames and --filter:
|
|
------------------------------------
|
|
author['id']
|
|
123
|
|
author['name']
|
|
test
|
|
author['self']
|
|
<circular reference>
|
|
category
|
|
test_category
|
|
extension
|
|
jpg
|
|
filename
|
|
1
|
|
num
|
|
1
|
|
subcategory
|
|
test_subcategory
|
|
tags[N]
|
|
0 foo
|
|
1 bar
|
|
2 テスト
|
|
user['id']
|
|
123
|
|
user['name']
|
|
test
|
|
user['self']
|
|
<circular reference>
|
|
""")
|
|
|
|
|
|
class TestUrlJob(TestJob):
|
|
jobclass = job.UrlJob
|
|
|
|
def test_default(self):
|
|
extr = TestExtractor.from_url("test:")
|
|
self.assertEqual(self._capture_stdout(extr), """\
|
|
https://example.org/1.jpg
|
|
https://example.org/2.jpg
|
|
https://example.org/3.jpg
|
|
""")
|
|
|
|
def test_fallback(self):
|
|
extr = TestExtractor.from_url("test:")
|
|
tjob = self.jobclass(extr)
|
|
tjob.handle_url = tjob.handle_url_fallback
|
|
|
|
self.assertEqual(self._capture_stdout(tjob), """\
|
|
https://example.org/1.jpg
|
|
| https://example.org/alt/1.jpg
|
|
https://example.org/2.jpg
|
|
| https://example.org/alt/2.jpg
|
|
https://example.org/3.jpg
|
|
| https://example.org/alt/3.jpg
|
|
""")
|
|
|
|
def test_parent(self):
|
|
extr = TestExtractorParent.from_url("test:parent")
|
|
self.assertEqual(self._capture_stdout(extr), """\
|
|
test:child
|
|
test:child
|
|
test:child
|
|
""")
|
|
|
|
def test_child(self):
|
|
extr = TestExtractorParent.from_url("test:parent")
|
|
tjob = job.UrlJob(extr, depth=0)
|
|
self.assertEqual(self._capture_stdout(tjob), 3 * """\
|
|
https://example.org/1.jpg
|
|
https://example.org/2.jpg
|
|
https://example.org/3.jpg
|
|
""")
|
|
|
|
|
|
class TestInfoJob(TestJob):
|
|
jobclass = job.InfoJob
|
|
|
|
def test_default(self):
|
|
extr = TestExtractor.from_url("test:")
|
|
self.assertEqual(self._capture_stdout(extr), """\
|
|
Category / Subcategory
|
|
"test_category" / "test_subcategory"
|
|
|
|
Filename format (default):
|
|
"test_{filename}.{extension}"
|
|
|
|
Directory format (default):
|
|
["{category}"]
|
|
|
|
""")
|
|
|
|
def test_custom(self):
|
|
config.set((), "filename", "custom")
|
|
config.set((), "directory", ("custom",))
|
|
config.set((), "sleep-request", 321)
|
|
extr = TestExtractor.from_url("test:")
|
|
extr.request_interval = 123.456
|
|
|
|
self.assertEqual(self._capture_stdout(extr), """\
|
|
Category / Subcategory
|
|
"test_category" / "test_subcategory"
|
|
|
|
Filename format (custom):
|
|
"custom"
|
|
Filename format (default):
|
|
"test_{filename}.{extension}"
|
|
|
|
Directory format (custom):
|
|
["custom"]
|
|
Directory format (default):
|
|
["{category}"]
|
|
|
|
Request interval (custom):
|
|
321
|
|
Request interval (default):
|
|
123.456
|
|
|
|
""")
|
|
|
|
def test_base_category(self):
|
|
extr = TestExtractor.from_url("test:")
|
|
extr.basecategory = "test_basecategory"
|
|
|
|
self.assertEqual(self._capture_stdout(extr), """\
|
|
Category / Subcategory / Basecategory
|
|
"test_category" / "test_subcategory" / "test_basecategory"
|
|
|
|
Filename format (default):
|
|
"test_{filename}.{extension}"
|
|
|
|
Directory format (default):
|
|
["{category}"]
|
|
|
|
""")
|
|
|
|
|
|
class TestDataJob(TestJob):
|
|
jobclass = job.DataJob
|
|
|
|
def test_default(self):
|
|
extr = TestExtractor.from_url("test:")
|
|
tjob = self.jobclass(extr, file=io.StringIO())
|
|
user = {"id": 123, "name": "test"}
|
|
|
|
tjob.run()
|
|
|
|
self.assertEqual(tjob.data, [
|
|
(Message.Directory, {
|
|
"category" : "test_category",
|
|
"subcategory": "test_subcategory",
|
|
"user" : user,
|
|
"author" : user,
|
|
}),
|
|
(Message.Url, "https://example.org/1.jpg", {
|
|
"category" : "test_category",
|
|
"subcategory": "test_subcategory",
|
|
"filename" : "1",
|
|
"extension" : "jpg",
|
|
"num" : 1,
|
|
"tags" : ["foo", "bar", "テスト"],
|
|
"user" : user,
|
|
"author" : user,
|
|
}),
|
|
(Message.Url, "https://example.org/2.jpg", {
|
|
"category" : "test_category",
|
|
"subcategory": "test_subcategory",
|
|
"filename" : "2",
|
|
"extension" : "jpg",
|
|
"num" : 2,
|
|
"tags" : ["foo", "bar", "テスト"],
|
|
"user" : user,
|
|
"author" : user,
|
|
}),
|
|
(Message.Url, "https://example.org/3.jpg", {
|
|
"category" : "test_category",
|
|
"subcategory": "test_subcategory",
|
|
"filename" : "3",
|
|
"extension" : "jpg",
|
|
"num" : 3,
|
|
"tags" : ["foo", "bar", "テスト"],
|
|
"user" : user,
|
|
"author" : user,
|
|
}),
|
|
])
|
|
|
|
def test_exception(self):
|
|
extr = TestExtractorException.from_url("test:exception")
|
|
tjob = self.jobclass(extr, file=io.StringIO())
|
|
tjob.run()
|
|
self.assertEqual(
|
|
tjob.data[-1], ("ZeroDivisionError", "division by zero"))
|
|
|
|
def test_private(self):
|
|
config.set(("output",), "private", True)
|
|
extr = TestExtractor.from_url("test:")
|
|
tjob = self.jobclass(extr, file=io.StringIO())
|
|
|
|
tjob.run()
|
|
|
|
for i in range(1, 4):
|
|
self.assertEqual(
|
|
tjob.data[i][2]["_fallback"],
|
|
("https://example.org/alt/{}.jpg".format(i),),
|
|
)
|
|
|
|
def test_sleep(self):
|
|
extr = TestExtractor.from_url("test:")
|
|
tjob = self.jobclass(extr, file=io.StringIO())
|
|
|
|
config.set((), "sleep-extractor", 123)
|
|
with patch("time.sleep") as sleep:
|
|
tjob.run()
|
|
sleep.assert_called_once_with(123)
|
|
|
|
config.set((), "sleep-extractor", 0)
|
|
with patch("time.sleep") as sleep:
|
|
tjob.run()
|
|
sleep.assert_not_called()
|
|
|
|
def test_ascii(self):
|
|
extr = TestExtractor.from_url("test:")
|
|
tjob = self.jobclass(extr)
|
|
|
|
tjob.file = buffer = io.StringIO()
|
|
tjob.run()
|
|
self.assertIn("""\
|
|
"tags": [
|
|
"foo",
|
|
"bar",
|
|
"\\u30c6\\u30b9\\u30c8"
|
|
],
|
|
""", buffer.getvalue())
|
|
|
|
tjob.file = buffer = io.StringIO()
|
|
tjob.ascii = False
|
|
tjob.run()
|
|
self.assertIn("""\
|
|
"tags": [
|
|
"foo",
|
|
"bar",
|
|
"テスト"
|
|
],
|
|
""", buffer.getvalue())
|
|
|
|
def test_num_string(self):
|
|
extr = TestExtractor.from_url("test:")
|
|
tjob = self.jobclass(extr, file=io.StringIO())
|
|
|
|
with patch("gallery_dl.util.number_to_string") as nts:
|
|
tjob.run()
|
|
self.assertEqual(len(nts.call_args_list), 0)
|
|
|
|
config.set(("output",), "num-to-str", True)
|
|
with patch("gallery_dl.util.number_to_string") as nts:
|
|
tjob.run()
|
|
self.assertEqual(len(nts.call_args_list), 72)
|
|
|
|
tjob.run()
|
|
self.assertEqual(tjob.data[-1][0], Message.Url)
|
|
self.assertEqual(tjob.data[-1][2]["num"], "3")
|
|
|
|
|
|
class TestExtractor(Extractor):
|
|
category = "test_category"
|
|
subcategory = "test_subcategory"
|
|
directory_fmt = ("{category}",)
|
|
filename_fmt = "test_{filename}.{extension}"
|
|
pattern = r"test:(child|self)?$"
|
|
|
|
def __init__(self, match):
|
|
Extractor.__init__(self, match)
|
|
self.user = {"id": 123, "name": "test"}
|
|
if match.group(1) == "self":
|
|
self.user["self"] = self.user
|
|
|
|
def items(self):
|
|
root = "https://example.org"
|
|
user = self.user
|
|
|
|
yield Message.Directory, {
|
|
"user": user,
|
|
"author": user,
|
|
}
|
|
|
|
for i in range(1, 4):
|
|
url = "{}/{}.jpg".format(root, i)
|
|
yield Message.Url, url, text.nameext_from_url(url, {
|
|
"num" : i,
|
|
"tags": ["foo", "bar", "テスト"],
|
|
"user": user,
|
|
"author": user,
|
|
"_fallback": ("{}/alt/{}.jpg".format(root, i),),
|
|
})
|
|
|
|
|
|
class TestExtractorParent(Extractor):
|
|
category = "test_category"
|
|
subcategory = "test_subcategory_parent"
|
|
pattern = r"test:parent"
|
|
|
|
def items(self):
|
|
url = "test:child"
|
|
|
|
for i in range(11, 14):
|
|
yield Message.Queue, url, {
|
|
"num" : i,
|
|
"tags": ["abc", "def"],
|
|
"_extractor": TestExtractor,
|
|
}
|
|
|
|
|
|
class TestExtractorException(Extractor):
|
|
category = "test_category"
|
|
subcategory = "test_subcategory_exception"
|
|
pattern = r"test:exception$"
|
|
|
|
def items(self):
|
|
return 1/0
|
|
|
|
|
|
class TestExtractorAlt(Extractor):
|
|
category = "test_category_alt"
|
|
subcategory = "test_subcategory"
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|