mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-10-03 00:35:32 -03:00
feat(downloads): return structured 429 rate-limit responses with retry_after
On a CivitAI/CivArchive 429, download-model and download-model-get now
return HTTP 429 with {"reason": "rate_limited", "retry_after": N}
instead of a generic 500 string, and the queue row goes back to
"queued" rather than history as failed — so queue drivers can
auto-pause and retry later instead of burning through the queue.
- new DownloadRateLimitError carrying retry_after/host (opt-in via
raise_on_rate_limit on Downloader; other call sites keep the legacy
string behavior)
- fail-fast pre-flight gate in DownloadManager consults
RateLimitCoordinator before acquiring the semaphore slot: hosts in
cooldown get an immediate structured 429, no HTTP request attempted
- best-effort 429 detection for the aria2 backend
This commit is contained in:
@@ -9,6 +9,8 @@ with ``success: false``, not 404. The browser extension's apiFetch treats any
|
||||
import json
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
from aiohttp import web
|
||||
@@ -184,3 +186,83 @@ async def test_retry_failed_history_returns_success(
|
||||
queue = await queue_service.get_queue()
|
||||
assert len(queue) == 1
|
||||
assert queue[0]["status"] == "queued"
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# Structured 429 rate-limit responses (download queue contract)
|
||||
# ----------------------------------------------------------------------
|
||||
|
||||
_RATE_LIMITED_RESULT = {
|
||||
"success": False,
|
||||
"reason": "rate_limited",
|
||||
"retry_after": 120,
|
||||
"error": "Download rate limited (429), retry after 120s",
|
||||
"download_id": "dl-1",
|
||||
}
|
||||
|
||||
|
||||
def _make_download_handler(result: dict) -> ModelDownloadHandler:
|
||||
return ModelDownloadHandler(
|
||||
ws_manager=None, # pyright: ignore[reportArgumentType] - unused by download endpoints
|
||||
logger=logging.getLogger("test-download-rate-limit"),
|
||||
download_use_case=SimpleNamespace(execute=AsyncMock(return_value=result)), # pyright: ignore[reportArgumentType]
|
||||
download_coordinator=None, # pyright: ignore[reportArgumentType] - unused by download endpoints
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_download_model_post_returns_429_when_rate_limited() -> None:
|
||||
"""POST /download-model maps reason="rate_limited" to HTTP 429."""
|
||||
handler = _make_download_handler(dict(_RATE_LIMITED_RESULT))
|
||||
request = SimpleNamespace(json=AsyncMock(return_value={"model_id": 1}))
|
||||
|
||||
response = await handler.download_model(request)
|
||||
|
||||
assert response.status == 429
|
||||
payload = json.loads(response.text)
|
||||
assert payload["success"] is False
|
||||
assert payload["reason"] == "rate_limited"
|
||||
assert payload["retry_after"] == 120
|
||||
assert payload["error"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_download_model_post_keeps_500_for_other_failures() -> None:
|
||||
handler = _make_download_handler({"success": False, "error": "boom"})
|
||||
request = SimpleNamespace(json=AsyncMock(return_value={"model_id": 1}))
|
||||
|
||||
response = await handler.download_model(request)
|
||||
|
||||
assert response.status == 500
|
||||
payload = json.loads(response.text)
|
||||
assert payload["success"] is False
|
||||
assert "reason" not in payload
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_download_model_get_returns_429_when_rate_limited() -> None:
|
||||
"""GET /download-model-get (extension queue driver) gets the same 429."""
|
||||
handler = _make_download_handler(dict(_RATE_LIMITED_RESULT))
|
||||
request = _queue_request("/api/lm/download-model-get", {"model_id": "1"})
|
||||
|
||||
response = await handler.download_model_get(request)
|
||||
|
||||
assert response.status == 429
|
||||
payload = json.loads(response.text)
|
||||
assert payload["success"] is False
|
||||
assert payload["reason"] == "rate_limited"
|
||||
assert payload["retry_after"] == 120
|
||||
assert payload["error"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_download_model_get_keeps_500_for_other_failures() -> None:
|
||||
handler = _make_download_handler({"success": False, "error": "boom"})
|
||||
request = _queue_request("/api/lm/download-model-get", {"model_id": "1"})
|
||||
|
||||
response = await handler.download_model_get(request)
|
||||
|
||||
assert response.status == 500
|
||||
payload = json.loads(response.text)
|
||||
assert payload["success"] is False
|
||||
assert "reason" not in payload
|
||||
|
||||
Reference in New Issue
Block a user