mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-10-09 02:52:12 -03:00
fix(download): scope the preflight cooldown gate to the download source
The pre-flight rate-limit gate checked civarchive.com for every download, so a cooldown armed by bulk metadata fetches blocked plain CivitAI downloads for the whole window. Only CivArchive-sourced downloads consult civarchive.com now.
This commit is contained in:
@@ -1969,7 +1969,7 @@ def _prepare_tracked_download(manager, download_id, status="downloading"):
|
||||
manager._pause_events[download_id] = DownloadStreamControl()
|
||||
|
||||
|
||||
async def _run_download(manager, download_id, tmp_path):
|
||||
async def _run_download(manager, download_id, tmp_path, source=None):
|
||||
return await manager._download_with_semaphore(
|
||||
download_id,
|
||||
1,
|
||||
@@ -1978,7 +1978,7 @@ async def _run_download(manager, download_id, tmp_path):
|
||||
"",
|
||||
None,
|
||||
False,
|
||||
None,
|
||||
source,
|
||||
None,
|
||||
False,
|
||||
)
|
||||
@@ -2063,6 +2063,59 @@ async def test_preflight_gate_blocks_when_host_in_cooldown(
|
||||
manager._download_semaphore.release()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_preflight_gate_ignores_civarchive_cooldown_for_civitai_downloads(
|
||||
monkeypatch, tmp_path, queue_service, reset_rate_limit_coordinator
|
||||
):
|
||||
"""A civarchive.com cooldown (armed e.g. by bulk metadata fetches) must
|
||||
not block a plain CivitAI download that never touches that host."""
|
||||
manager = DownloadManager()
|
||||
monkeypatch.setattr(manager, "_cleanup_download_record", AsyncMock())
|
||||
|
||||
coordinator = await RateLimitCoordinator.get_instance()
|
||||
coordinator.register_rate_limit("civarchive.com", 1800)
|
||||
|
||||
download_id = "dl-civarchive-cooldown"
|
||||
await queue_service.add_to_queue(download_id=download_id, model_id=1)
|
||||
_prepare_tracked_download(manager, download_id, status="waiting")
|
||||
|
||||
execute = AsyncMock(return_value={"success": True})
|
||||
monkeypatch.setattr(manager, "_execute_original_download", execute)
|
||||
|
||||
result = await _run_download(manager, download_id, tmp_path)
|
||||
|
||||
assert execute.await_count == 1
|
||||
assert result["success"] is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_preflight_gate_blocks_civarchive_source_during_civarchive_cooldown(
|
||||
monkeypatch, tmp_path, queue_service, reset_rate_limit_coordinator
|
||||
):
|
||||
"""Downloads explicitly sourced from CivArchive still gate on its cooldown."""
|
||||
manager = DownloadManager()
|
||||
monkeypatch.setattr(manager, "_cleanup_download_record", AsyncMock())
|
||||
|
||||
coordinator = await RateLimitCoordinator.get_instance()
|
||||
coordinator.register_rate_limit("civarchive.com", 1800)
|
||||
|
||||
download_id = "dl-civarchive-source"
|
||||
await queue_service.add_to_queue(download_id=download_id, model_id=1)
|
||||
_prepare_tracked_download(manager, download_id, status="waiting")
|
||||
|
||||
execute = AsyncMock(
|
||||
side_effect=AssertionError("download must not start during cooldown")
|
||||
)
|
||||
monkeypatch.setattr(manager, "_execute_original_download", execute)
|
||||
|
||||
result = await _run_download(manager, download_id, tmp_path, source="civarchive")
|
||||
|
||||
assert execute.await_count == 0
|
||||
assert result["success"] is False
|
||||
assert result["reason"] == "rate_limited"
|
||||
assert "civarchive.com" in result["error"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rate_limit_retry_after_falls_back_to_coordinator_backoff(
|
||||
monkeypatch, tmp_path, queue_service, reset_rate_limit_coordinator
|
||||
|
||||
Reference in New Issue
Block a user