Dariusz L
f7dfe3d8e5
- Recover completed .part files on HTTP 416 without deleting and redownloading - Add stream stall detection and automatic reconnect for interrupted socket reads - Implement persistent download state store and pause/resume lifecycle controls - Add RateLimitCoordinator and ConnectivityGuard circuit breaker for requests - Support proxy_url configuration and validate aria2 proxy settings - Add unit tests for 416 recovery, stall reconnect, state store, and network guard
212 lines
7.0 KiB
Python
212 lines
7.0 KiB
Python
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
import requests
|
|
|
|
from core.network_guard import (
|
|
ConnectivityGuard,
|
|
RateLimitCoordinator,
|
|
get_connectivity_guard,
|
|
get_rate_limit_coordinator,
|
|
normalize_destination,
|
|
parse_retry_after,
|
|
)
|
|
from core.network_utils import request_source_response
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def clean_guards():
|
|
"""Ensure both guards start each test with clean state."""
|
|
rl = get_rate_limit_coordinator()
|
|
cg = get_connectivity_guard()
|
|
rl.clear()
|
|
cg.clear()
|
|
yield
|
|
rl.clear()
|
|
cg.clear()
|
|
|
|
|
|
def test_normalize_destination():
|
|
assert normalize_destination("https://civitai.com/api/v1/models") == "civitai.com"
|
|
assert normalize_destination("http://huggingface.co:443/models/test") == "huggingface.co"
|
|
assert normalize_destination("sub.domain.org/path?query=1") == "sub.domain.org"
|
|
assert normalize_destination("api.example.com") == "api.example.com"
|
|
assert normalize_destination("") == "__global__"
|
|
assert normalize_destination(None) == "__global__"
|
|
|
|
|
|
def test_parse_retry_after():
|
|
assert parse_retry_after("120") == 120
|
|
assert parse_retry_after(45) == 45
|
|
assert parse_retry_after(None, default=30) == 30
|
|
assert parse_retry_after("invalid", default=60) == 60
|
|
assert parse_retry_after("-10") == 1 # clamped min 1
|
|
assert parse_retry_after("10000") == 3600 # clamped max 3600
|
|
|
|
|
|
def test_rate_limit_coordinator_lifecycle():
|
|
coordinator = RateLimitCoordinator()
|
|
url = "https://civitai.com/api/download/models/123"
|
|
|
|
assert not coordinator.is_rate_limited(url)
|
|
assert coordinator.get_cooldown(url) == 0.0
|
|
|
|
registered = coordinator.register_rate_limit(url, 2, reason="Rate limit test")
|
|
assert registered == 2
|
|
assert coordinator.is_rate_limited(url)
|
|
assert coordinator.is_rate_limited("https://civitai.com/other") # same domain
|
|
assert not coordinator.is_rate_limited("https://huggingface.co/model") # different domain
|
|
assert coordinator.get_reason(url) == "Rate limit test"
|
|
|
|
cooldown = coordinator.get_cooldown(url)
|
|
assert 0.0 < cooldown <= 2.0
|
|
|
|
# Test clearing
|
|
coordinator.clear(url)
|
|
assert not coordinator.is_rate_limited(url)
|
|
assert coordinator.get_cooldown(url) == 0.0
|
|
|
|
|
|
def test_rate_limit_coordinator_wait_if_needed():
|
|
coordinator = RateLimitCoordinator()
|
|
url = "https://example.com/test"
|
|
|
|
# Not rate limited -> returns True immediately
|
|
assert coordinator.wait_if_needed(url, max_wait=1.0) is True
|
|
|
|
# Active rate limit longer than max_wait -> returns False immediately without sleeping
|
|
coordinator.register_rate_limit(url, 10)
|
|
assert coordinator.wait_if_needed(url, max_wait=0.1) is False
|
|
|
|
|
|
def test_connectivity_guard_failure_and_recovery():
|
|
guard = ConnectivityGuard(failure_threshold=3, base_backoff=5.0)
|
|
url = "https://civitai.com/api/test"
|
|
|
|
assert not guard.is_offline(url)
|
|
|
|
# Record 2 failures (under threshold of 3)
|
|
exc = requests.exceptions.ConnectionError("Connection refused")
|
|
assert guard.record_failure(url, exc) is False
|
|
assert not guard.is_offline(url)
|
|
|
|
assert guard.record_failure(url, exc) is False
|
|
assert not guard.is_offline(url)
|
|
|
|
# 3rd failure activates guard
|
|
assert guard.record_failure(url, exc) is True
|
|
assert guard.is_offline(url)
|
|
assert guard.get_offline_cooldown(url) > 0.0
|
|
|
|
# Non-transport error (e.g. HTTP 404) should not increment transport failures
|
|
http_err = requests.exceptions.HTTPError("404 Not Found")
|
|
assert guard.record_failure(url, http_err) is False
|
|
|
|
# Success resets the guard
|
|
guard.record_success(url)
|
|
assert not guard.is_offline(url)
|
|
assert guard.get_offline_cooldown(url) == 0.0
|
|
|
|
|
|
def test_request_source_response_with_rate_limit_cooldown():
|
|
url = "https://api.civitai.com/v1/models"
|
|
coordinator = get_rate_limit_coordinator()
|
|
|
|
# Pre-register a 10s cooldown
|
|
coordinator.register_rate_limit(url, 10, reason="CivitAI rate limited")
|
|
|
|
# request_source_response should skip immediately
|
|
with patch("requests.get") as mock_get:
|
|
response = request_source_response(url, log_name="Test")
|
|
assert response is None
|
|
mock_get.assert_not_called()
|
|
|
|
# If raise_on_error is True, raises HTTPError
|
|
with pytest.raises(requests.exceptions.HTTPError, match="Rate limit cooldown active"):
|
|
request_source_response(url, log_name="Test", raise_on_error=True)
|
|
|
|
|
|
def test_request_source_response_registers_429():
|
|
url = "https://api.civitai.com/v1/models"
|
|
coordinator = get_rate_limit_coordinator()
|
|
assert not coordinator.is_rate_limited(url)
|
|
|
|
mock_resp = MagicMock()
|
|
mock_resp.status_code = 429
|
|
mock_resp.headers = {"Retry-After": "30"}
|
|
|
|
with patch("requests.get", return_value=mock_resp), patch("time.sleep"):
|
|
resp = request_source_response(url, max_attempts=1)
|
|
assert resp == mock_resp
|
|
|
|
assert coordinator.is_rate_limited(url)
|
|
assert 20.0 < coordinator.get_cooldown(url) <= 30.0
|
|
|
|
|
|
def test_request_source_response_connectivity_guard_skip():
|
|
url = "https://offline.service.org/api"
|
|
guard = get_connectivity_guard()
|
|
|
|
# Activate offline guard
|
|
exc = requests.exceptions.ConnectionError("Network is unreachable")
|
|
for _ in range(guard.failure_threshold):
|
|
guard.record_failure(url, exc)
|
|
|
|
assert guard.is_offline(url)
|
|
|
|
# Request should be skipped without hitting requests.get
|
|
with patch("requests.get") as mock_get:
|
|
resp = request_source_response(url)
|
|
assert resp is None
|
|
mock_get.assert_not_called()
|
|
|
|
|
|
def test_download_file_defers_when_rate_limited(tmp_path):
|
|
from core.download.api import context as downloader
|
|
from core.download.orchestrator import download_file
|
|
|
|
url = "https://civitai.com/api/download/models/999"
|
|
dest_path = str(tmp_path / "rate_limited_model.safetensors")
|
|
coordinator = get_rate_limit_coordinator()
|
|
|
|
# Pre-register 60s cooldown
|
|
coordinator.register_rate_limit(url, 60, reason="HTTP 429 civitai")
|
|
|
|
result = download_file(
|
|
url,
|
|
dest_path,
|
|
"dl-rl-test",
|
|
dependencies=downloader,
|
|
)
|
|
|
|
assert result["success"] is False
|
|
assert "Rate limit exceeded (HTTP 429)" in result["error"]
|
|
assert downloader.download_progress["dl-rl-test"]["status"] == "error"
|
|
|
|
|
|
def test_download_file_defers_when_connectivity_guard_offline(tmp_path):
|
|
from core.download.api import context as downloader
|
|
from core.download.orchestrator import download_file
|
|
|
|
url = "https://offline-cdn.example.org/models/test.safetensors"
|
|
dest_path = str(tmp_path / "offline_model.safetensors")
|
|
guard = get_connectivity_guard()
|
|
|
|
exc = requests.exceptions.ConnectionError("Connection refused")
|
|
for _ in range(guard.failure_threshold):
|
|
guard.record_failure(url, exc)
|
|
|
|
assert guard.is_offline(url)
|
|
|
|
result = download_file(
|
|
url,
|
|
dest_path,
|
|
"dl-cg-test",
|
|
dependencies=downloader,
|
|
)
|
|
|
|
assert result["success"] is False
|
|
assert "Network offline" in result["error"]
|
|
assert downloader.download_progress["dl-cg-test"]["status"] == "error"
|
|
|