mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-05-06 16:36:45 -03:00
Merge pull request #899 from Phinease/fix/resumable-download-retries
fix: preserve resumable downloads across retries
This commit is contained in:
@@ -30,10 +30,21 @@ class FakeStream:
|
||||
|
||||
|
||||
class FakeResponse:
|
||||
def __init__(self, status, headers, chunks):
|
||||
def __init__(
|
||||
self,
|
||||
status,
|
||||
headers,
|
||||
chunks,
|
||||
*,
|
||||
url="https://example.com/file",
|
||||
history=None,
|
||||
):
|
||||
self.status = status
|
||||
self.headers = headers
|
||||
self.content = FakeStream(chunks)
|
||||
self.url = url
|
||||
self.history = history or []
|
||||
self.released = False
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
@@ -41,14 +52,25 @@ class FakeResponse:
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
def release(self):
|
||||
self.released = True
|
||||
|
||||
|
||||
class FakeSession:
|
||||
def __init__(self, responses):
|
||||
self._responses = list(responses)
|
||||
self._get_calls = 0
|
||||
self.requests = []
|
||||
|
||||
def get(self, url, headers=None, allow_redirects=True, proxy=None): # noqa: D401 - signature mirrors aiohttp
|
||||
del url, headers, allow_redirects, proxy
|
||||
self.requests.append(
|
||||
{
|
||||
"url": url,
|
||||
"headers": headers or {},
|
||||
"allow_redirects": allow_redirects,
|
||||
"proxy": proxy,
|
||||
}
|
||||
)
|
||||
response_factory = self._responses[self._get_calls]
|
||||
self._get_calls += 1
|
||||
return response_factory()
|
||||
@@ -75,7 +97,7 @@ def _build_downloader(responses, *, max_retries=0):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_download_file_fails_when_size_mismatch(tmp_path):
|
||||
async def test_download_file_preserves_incomplete_part_when_size_mismatch(tmp_path):
|
||||
target_path = tmp_path / "model" / "file.bin"
|
||||
target_path.parent.mkdir()
|
||||
|
||||
@@ -94,7 +116,7 @@ async def test_download_file_fails_when_size_mismatch(tmp_path):
|
||||
assert success is False
|
||||
assert "mismatch" in message.lower()
|
||||
assert not target_path.exists()
|
||||
assert not Path(str(target_path) + ".part").exists()
|
||||
assert Path(str(target_path) + ".part").read_bytes() == b"abc"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -136,7 +158,9 @@ async def test_download_file_succeeds_when_sizes_match(tmp_path):
|
||||
|
||||
downloader = _build_downloader(responses)
|
||||
|
||||
success, result_path = await downloader.download_file("https://example.com/file", str(target_path))
|
||||
success, result_path = await downloader.download_file(
|
||||
"https://example.com/file", str(target_path)
|
||||
)
|
||||
|
||||
assert success is True
|
||||
assert Path(result_path).read_bytes() == payload
|
||||
@@ -166,9 +190,77 @@ async def test_download_file_recovers_from_stall(tmp_path):
|
||||
downloader = _build_downloader(responses, max_retries=1)
|
||||
downloader.stall_timeout = 0.05
|
||||
|
||||
success, result_path = await downloader.download_file("https://example.com/file", str(target_path))
|
||||
success, result_path = await downloader.download_file(
|
||||
"https://example.com/file", str(target_path)
|
||||
)
|
||||
|
||||
assert success is True
|
||||
assert Path(result_path).read_bytes() == payload
|
||||
assert downloader._session._get_calls == 2
|
||||
assert not Path(str(target_path) + ".part").exists()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_download_file_resumes_after_incomplete_integrity_check(tmp_path):
|
||||
target_path = tmp_path / "model" / "file.bin"
|
||||
target_path.parent.mkdir()
|
||||
|
||||
responses = [
|
||||
lambda: FakeResponse(
|
||||
status=200,
|
||||
headers={"content-length": "6"},
|
||||
chunks=[b"abc"],
|
||||
),
|
||||
lambda: FakeResponse(
|
||||
status=206,
|
||||
headers={"content-length": "3", "Content-Range": "bytes 3-5/6"},
|
||||
chunks=[b"def"],
|
||||
),
|
||||
]
|
||||
|
||||
downloader = _build_downloader(responses, max_retries=1)
|
||||
|
||||
success, result_path = await downloader.download_file("https://example.com/file", str(target_path))
|
||||
|
||||
assert success is True
|
||||
assert Path(result_path).read_bytes() == b"abcdef"
|
||||
assert downloader._session._get_calls == 2
|
||||
assert downloader._session.requests[1]["headers"]["Range"] == "bytes=3-"
|
||||
assert not Path(str(target_path) + ".part").exists()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_download_file_retries_redirected_url_when_range_not_honored(tmp_path):
|
||||
target_path = tmp_path / "model" / "file.bin"
|
||||
target_path.parent.mkdir()
|
||||
Path(str(target_path) + ".part").write_bytes(b"abc")
|
||||
|
||||
redirected_url = "https://download.example.com/file.bin"
|
||||
first_response = FakeResponse(
|
||||
status=200,
|
||||
headers={"content-length": "6"},
|
||||
chunks=[],
|
||||
url=redirected_url,
|
||||
history=[object()],
|
||||
)
|
||||
|
||||
responses = [
|
||||
lambda: first_response,
|
||||
lambda: FakeResponse(
|
||||
status=206,
|
||||
headers={"content-length": "3", "Content-Range": "bytes 3-5/6"},
|
||||
chunks=[b"def"],
|
||||
url=redirected_url,
|
||||
),
|
||||
]
|
||||
|
||||
downloader = _build_downloader(responses, max_retries=0)
|
||||
|
||||
success, result_path = await downloader.download_file("https://example.com/file", str(target_path))
|
||||
|
||||
assert success is True
|
||||
assert Path(result_path).read_bytes() == b"abcdef"
|
||||
assert first_response.released is True
|
||||
assert downloader._session.requests[0]["headers"]["Range"] == "bytes=3-"
|
||||
assert downloader._session.requests[1]["url"] == redirected_url
|
||||
assert downloader._session.requests[1]["headers"]["Range"] == "bytes=3-"
|
||||
|
||||
Reference in New Issue
Block a user