fix: respect sleep-requests arg in specific code paths.
This commit is contained in:
parent
ac3c666d32
commit
9b3cc8b310
9 changed files with 149 additions and 36 deletions
|
|
@ -187,6 +187,7 @@ class GenericTaskHandler(BaseHandler):
|
|||
url=url,
|
||||
no_archive=True,
|
||||
no_log=True,
|
||||
budget_sleep=True,
|
||||
)
|
||||
|
||||
if not info:
|
||||
|
|
|
|||
|
|
@ -185,6 +185,7 @@ class RssGenericHandler(BaseHandler):
|
|||
url=url,
|
||||
no_archive=True,
|
||||
no_log=True,
|
||||
budget_sleep=True,
|
||||
)
|
||||
|
||||
if not info:
|
||||
|
|
|
|||
|
|
@ -107,6 +107,7 @@ class HandleTask(TaskSchema):
|
|||
no_archive=True,
|
||||
follow_redirect=False,
|
||||
sanitize_info=True,
|
||||
budget_sleep=True,
|
||||
)
|
||||
|
||||
if not ie_info or not isinstance(ie_info, dict):
|
||||
|
|
@ -133,7 +134,7 @@ class HandleTask(TaskSchema):
|
|||
|
||||
archive_file: Path = Path(archive_file)
|
||||
|
||||
(ie_info, _) = await fetch_info(params, self.url, no_archive=True, follow_redirect=True)
|
||||
(ie_info, _) = await fetch_info(params, self.url, no_archive=True, follow_redirect=True, budget_sleep=True)
|
||||
if not ie_info or not isinstance(ie_info, dict):
|
||||
return (False, "Failed to extract information from URL.")
|
||||
|
||||
|
|
|
|||
|
|
@ -49,6 +49,17 @@ class ExtractorConfig:
|
|||
self.wait_threshold = wait_threshold
|
||||
|
||||
|
||||
def _sleep_timeout(config: dict[str, Any], timeout: float, budget_sleep: bool) -> float:
|
||||
if not budget_sleep:
|
||||
return timeout
|
||||
|
||||
sleep_requests = config.get("sleep_interval_requests")
|
||||
if not isinstance(sleep_requests, int | float) or sleep_requests <= 0:
|
||||
return timeout
|
||||
|
||||
return timeout + min(float(sleep_requests) * 20, 300.0)
|
||||
|
||||
|
||||
class ExtractorPool(metaclass=Singleton):
|
||||
"""
|
||||
Manages process pool and semaphore for video information extraction.
|
||||
|
|
@ -312,6 +323,7 @@ async def fetch_info(
|
|||
sanitize_info: bool = False,
|
||||
capture_logs: int | None = None,
|
||||
extractor_config: ExtractorConfig | None = None,
|
||||
budget_sleep: bool = False,
|
||||
**kwargs,
|
||||
) -> tuple[dict[str, Any] | None, list[dict[str, Any]]]:
|
||||
"""
|
||||
|
|
@ -329,6 +341,7 @@ async def fetch_info(
|
|||
sanitize_info: Sanitize the extracted information
|
||||
capture_logs: If provided (e.g., logging.WARNING), capture logs
|
||||
extractor_config: Configuration for the extractor
|
||||
budget_sleep: Whether to add extra timeout budget for request-sleep-heavy extraction
|
||||
**kwargs: Additional arguments
|
||||
|
||||
Returns:
|
||||
|
|
@ -352,6 +365,7 @@ async def fetch_info(
|
|||
loop = asyncio.get_running_loop()
|
||||
|
||||
safe_config = _sanitize_config(config)
|
||||
timeout = _sleep_timeout(safe_config, extractor_config.timeout, budget_sleep)
|
||||
|
||||
try:
|
||||
try:
|
||||
|
|
@ -372,7 +386,7 @@ async def fetch_info(
|
|||
**kwargs,
|
||||
),
|
||||
),
|
||||
timeout=extractor_config.timeout,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
except TimeoutError:
|
||||
|
|
@ -396,7 +410,7 @@ async def fetch_info(
|
|||
**kwargs,
|
||||
),
|
||||
),
|
||||
timeout=extractor_config.timeout,
|
||||
timeout=timeout,
|
||||
)
|
||||
finally:
|
||||
semaphore.release()
|
||||
|
|
|
|||
|
|
@ -208,6 +208,7 @@ async def add(
|
|||
no_archive=False,
|
||||
follow_redirect=True,
|
||||
capture_logs=logging.WARNING,
|
||||
budget_sleep=True,
|
||||
)
|
||||
|
||||
if not entry:
|
||||
|
|
|
|||
|
|
@ -792,6 +792,7 @@ async def item_nfo_generate(request: Request, queue: DownloadQueue) -> Response:
|
|||
url=item.info.url,
|
||||
no_archive=True,
|
||||
follow_redirect=True,
|
||||
budget_sleep=True,
|
||||
)
|
||||
|
||||
if not info_dict:
|
||||
|
|
|
|||
|
|
@ -18,7 +18,7 @@ def reset_config() -> Generator[None, None, None]:
|
|||
Config._reset_singleton()
|
||||
|
||||
|
||||
class _JsonRequest:
|
||||
class _Req:
|
||||
def __init__(self, payload: Any) -> None:
|
||||
self._payload = payload
|
||||
self.body_exists = payload is not None
|
||||
|
|
@ -27,65 +27,65 @@ class _JsonRequest:
|
|||
return self._payload
|
||||
|
||||
|
||||
class _InspectRequest(_JsonRequest):
|
||||
class _InspectReq(_Req):
|
||||
query: dict[str, str] = {}
|
||||
match_info: dict[str, str] = {}
|
||||
|
||||
|
||||
class _QueryRequest:
|
||||
class _QueryReq:
|
||||
def __init__(self, query: dict[str, str]) -> None:
|
||||
self.query = query
|
||||
|
||||
|
||||
def _patch_to_thread(monkeypatch: pytest.MonkeyPatch, module: Any, config: Config, url: str) -> dict[str, bool]:
|
||||
called = {"to_thread": False, "validate": False}
|
||||
def _patch_thread(monkeypatch: pytest.MonkeyPatch, module: Any, config: Config, url: str) -> dict[str, bool]:
|
||||
seen = {"to_thread": False, "validate": False}
|
||||
|
||||
def fake_validate_url(next_url: str, allow_internal: bool = False) -> bool:
|
||||
called["validate"] = True
|
||||
seen["validate"] = True
|
||||
assert next_url == url
|
||||
assert allow_internal is config.allow_internal_urls
|
||||
raise ValueError("Invalid hostname.")
|
||||
|
||||
async def fake_to_thread(func, *args, **kwargs):
|
||||
called["to_thread"] = True
|
||||
seen["to_thread"] = True
|
||||
return func(*args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(module, "validate_url", fake_validate_url)
|
||||
monkeypatch.setattr(module.asyncio, "to_thread", fake_to_thread)
|
||||
return called
|
||||
return seen
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_inspect_validate_off_thread(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
async def test_inspect_thread(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
config = Config.get_instance()
|
||||
request = _InspectRequest({"url": "https://bad.example/task"})
|
||||
called = _patch_to_thread(monkeypatch, tasks_router, config, "https://bad.example/task")
|
||||
request = _InspectReq({"url": "https://bad.example/task"})
|
||||
seen = _patch_thread(monkeypatch, tasks_router, config, "https://bad.example/task")
|
||||
|
||||
response = await tasks_router.task_handler_inspect(request, handler=None, encoder=None, config=config)
|
||||
|
||||
assert response.status == 400
|
||||
assert json.loads(response.body.decode("utf-8")) == {"error": "Invalid hostname."}
|
||||
assert called == {"to_thread": True, "validate": True}
|
||||
assert seen == {"to_thread": True, "validate": True}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_conditions_validate_off_thread(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
async def test_conditions_thread(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
config = Config.get_instance()
|
||||
request = _JsonRequest({"url": "https://bad.example/cond", "condition": "title ~= 'x'"})
|
||||
called = _patch_to_thread(monkeypatch, conditions_router, config, "https://bad.example/cond")
|
||||
request = _Req({"url": "https://bad.example/cond", "condition": "title ~= 'x'"})
|
||||
seen = _patch_thread(monkeypatch, conditions_router, config, "https://bad.example/cond")
|
||||
|
||||
response = await conditions_router.conditions_test(request, encoder=None, cache=None, config=config)
|
||||
|
||||
assert response.status == 400
|
||||
assert json.loads(response.body.decode("utf-8")) == {"error": "Invalid hostname."}
|
||||
assert called == {"to_thread": True, "validate": True}
|
||||
assert seen == {"to_thread": True, "validate": True}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_info_validate_off_thread(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
async def test_info_thread(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
config = Config.get_instance()
|
||||
request = _QueryRequest({"url": "https://bad.example/info"})
|
||||
called = _patch_to_thread(monkeypatch, ytdlp_router, config, "https://bad.example/info")
|
||||
request = _QueryReq({"url": "https://bad.example/info"})
|
||||
seen = _patch_thread(monkeypatch, ytdlp_router, config, "https://bad.example/info")
|
||||
|
||||
response = await ytdlp_router.get_info(request, cache=None, config=config)
|
||||
|
||||
|
|
@ -95,14 +95,14 @@ async def test_info_validate_off_thread(monkeypatch: pytest.MonkeyPatch) -> None
|
|||
"message": "Invalid hostname.",
|
||||
"error": "Invalid hostname.",
|
||||
}
|
||||
assert called == {"to_thread": True, "validate": True}
|
||||
assert seen == {"to_thread": True, "validate": True}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_archive_ids_validate_off_thread(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
async def test_archive_ids_thread(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
config = Config.get_instance()
|
||||
request = _JsonRequest(["https://bad.example/archive"])
|
||||
called = _patch_to_thread(monkeypatch, ytdlp_router, config, "https://bad.example/archive")
|
||||
request = _Req(["https://bad.example/archive"])
|
||||
seen = _patch_thread(monkeypatch, ytdlp_router, config, "https://bad.example/archive")
|
||||
|
||||
response = await ytdlp_router.get_archive_ids(request, config)
|
||||
|
||||
|
|
@ -117,4 +117,4 @@ async def test_archive_ids_validate_off_thread(monkeypatch: pytest.MonkeyPatch)
|
|||
"error": "Invalid hostname.",
|
||||
}
|
||||
]
|
||||
assert called == {"to_thread": True, "validate": True}
|
||||
assert seen == {"to_thread": True, "validate": True}
|
||||
|
|
|
|||
94
app/tests/test_extractor_timeout.py
Normal file
94
app/tests/test_extractor_timeout.py
Normal file
|
|
@ -0,0 +1,94 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
|
||||
from app.features.ytdlp import extractor
|
||||
|
||||
|
||||
class _Loop:
|
||||
def __init__(self) -> None:
|
||||
self.calls: list[object | None] = []
|
||||
|
||||
def run_in_executor(self, executor, func): # noqa: ANN001
|
||||
self.calls.append(executor)
|
||||
return func
|
||||
|
||||
|
||||
class _Pool:
|
||||
def __init__(self) -> None:
|
||||
self.semaphore = asyncio.Semaphore(1)
|
||||
self.executor = object()
|
||||
|
||||
def get_semaphore(self, _config: extractor.ExtractorConfig) -> asyncio.Semaphore:
|
||||
return self.semaphore
|
||||
|
||||
def get_pool(self, _config: extractor.ExtractorConfig) -> object:
|
||||
return self.executor
|
||||
|
||||
|
||||
def test_sleep_budget() -> None:
|
||||
assert extractor._sleep_timeout({}, 70, False) == 70
|
||||
assert extractor._sleep_timeout({"sleep_interval_requests": 0}, 70, True) == 70
|
||||
assert extractor._sleep_timeout({"sleep_interval_requests": 3}, 70, True) == 130
|
||||
assert extractor._sleep_timeout({"sleep_interval_requests": 30}, 70, True) == 370
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_timeout_no_retry(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
pool = _Pool()
|
||||
loop = _Loop()
|
||||
seen: list[float] = []
|
||||
|
||||
async def fake_wait_for(*, fut, timeout):
|
||||
seen.append(timeout)
|
||||
raise TimeoutError
|
||||
|
||||
monkeypatch.setattr(extractor.ExtractorPool, "get_instance", classmethod(lambda cls: pool))
|
||||
monkeypatch.setattr(extractor.asyncio, "get_running_loop", lambda: loop)
|
||||
monkeypatch.setattr(extractor.asyncio, "wait_for", fake_wait_for)
|
||||
|
||||
with pytest.raises(TimeoutError):
|
||||
await extractor.fetch_info(
|
||||
config={"sleep_interval_requests": 3},
|
||||
url="https://example.com",
|
||||
extractor_config=extractor.ExtractorConfig(concurrency=1, timeout=70),
|
||||
budget_sleep=True,
|
||||
)
|
||||
|
||||
assert loop.calls == [pool.executor]
|
||||
assert seen == [130]
|
||||
assert not pool.semaphore.locked()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pool_fallback(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
pool = _Pool()
|
||||
loop = _Loop()
|
||||
expected = ({"id": "ok"}, [])
|
||||
seen: list[float] = []
|
||||
|
||||
async def fake_wait_for(*, fut, timeout):
|
||||
seen.append(timeout)
|
||||
if len(loop.calls) == 1:
|
||||
raise RuntimeError("pool failed")
|
||||
return fut()
|
||||
|
||||
monkeypatch.setattr(extractor.ExtractorPool, "get_instance", classmethod(lambda cls: pool))
|
||||
monkeypatch.setattr(extractor.asyncio, "get_running_loop", lambda: loop)
|
||||
monkeypatch.setattr(extractor.asyncio, "wait_for", fake_wait_for)
|
||||
monkeypatch.setattr(extractor, "extract_info_sync", Mock(return_value=expected))
|
||||
|
||||
result = await extractor.fetch_info(
|
||||
config={"sleep_interval_requests": 3},
|
||||
url="https://example.com",
|
||||
extractor_config=extractor.ExtractorConfig(concurrency=1, timeout=70),
|
||||
budget_sleep=True,
|
||||
)
|
||||
|
||||
assert result == expected
|
||||
assert loop.calls == [pool.executor, None]
|
||||
assert seen == [130, 130]
|
||||
assert not pool.semaphore.locked()
|
||||
|
|
@ -18,7 +18,7 @@ def reset_config() -> Generator[None, None, None]:
|
|||
Config._reset_singleton()
|
||||
|
||||
|
||||
class _Response:
|
||||
class _Resp:
|
||||
def __init__(self, *, status_code: int = 200, content: bytes = b"img", content_type: str = "image/jpeg") -> None:
|
||||
self.status_code = status_code
|
||||
self.content = content
|
||||
|
|
@ -26,25 +26,25 @@ class _Response:
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_thumbnail_validate_off_thread(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
async def test_thumb_thread(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
config = Config.get_instance()
|
||||
req = make_mocked_request("GET", "/api/thumbnail?url=https://example.com/a.jpg")
|
||||
req._rel_url = req._rel_url.with_query({"url": "https://example.com/a.jpg"})
|
||||
|
||||
called = {"to_thread": False, "validate": False}
|
||||
seen = {"to_thread": False, "validate": False}
|
||||
|
||||
def fake_validate_url(url: str, allow_internal: bool = False) -> bool:
|
||||
called["validate"] = True
|
||||
seen["validate"] = True
|
||||
assert url == "https://example.com/a.jpg"
|
||||
assert allow_internal is config.allow_internal_urls
|
||||
return True
|
||||
|
||||
async def fake_to_thread(func, *args, **kwargs):
|
||||
called["to_thread"] = True
|
||||
seen["to_thread"] = True
|
||||
return func(*args, **kwargs)
|
||||
|
||||
client = AsyncMock()
|
||||
client.request.return_value = _Response()
|
||||
client.request.return_value = _Resp()
|
||||
|
||||
monkeypatch.setattr(images, "validate_url", fake_validate_url)
|
||||
monkeypatch.setattr(images.asyncio, "to_thread", fake_to_thread)
|
||||
|
|
@ -70,13 +70,13 @@ async def test_thumbnail_validate_off_thread(monkeypatch: pytest.MonkeyPatch) ->
|
|||
response = await images.get_thumbnail(req, config)
|
||||
|
||||
assert response.status == web.HTTPOk.status_code
|
||||
assert called["to_thread"] is True
|
||||
assert called["validate"] is True
|
||||
assert seen["to_thread"] is True
|
||||
assert seen["validate"] is True
|
||||
client.request.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_thumbnail_validate_rejects(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
async def test_thumb_reject(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
config = Config.get_instance()
|
||||
req = make_mocked_request("GET", "/api/thumbnail?url=https://bad.example/a.jpg")
|
||||
req._rel_url = req._rel_url.with_query({"url": "https://bad.example/a.jpg"})
|
||||
|
|
|
|||
Loading…
Reference in a new issue