fix(tests): update tests to match current download implementation

- Remove calculate_sha256 mocking from download_manager tests since
  SHA256 now comes from API metadata (not recalculated during download)
- Update chunk_size assertion from 4MB to 16MB in downloader config test
This commit is contained in:
Will Miao
2026-03-26 18:00:04 +08:00
parent 95e5bc26d1
commit 3b001a6cd8
2 changed files with 110 additions and 102 deletions

View File

@@ -281,8 +281,6 @@ async def test_execute_download_extracts_zip_single_model(monkeypatch, tmp_path)
DownloadManager, "_get_lora_scanner", AsyncMock(return_value=dummy_scanner)
)
monkeypatch.setattr(MetadataManager, "save_metadata", AsyncMock(return_value=True))
hash_calculator = AsyncMock(return_value="hash-single")
monkeypatch.setattr(download_manager, "calculate_sha256", hash_calculator)
result = await manager._execute_download(
download_urls=download_urls,
@@ -299,10 +297,10 @@ async def test_execute_download_extracts_zip_single_model(monkeypatch, tmp_path)
assert not zip_path.exists()
extracted = save_dir / "model.safetensors"
assert extracted.exists()
assert hash_calculator.await_args.args[0] == str(extracted)
saved_call = MetadataManager.save_metadata.await_args
assert saved_call.args[0] == str(extracted)
assert saved_call.args[1].sha256 == "hash-single"
# SHA256 comes from metadata (API value), not recalculated
assert saved_call.args[1].sha256 == "sha256"
assert dummy_scanner.add_model_to_cache.await_count == 1
@@ -351,8 +349,6 @@ async def test_execute_download_extracts_zip_multiple_models(monkeypatch, tmp_pa
DownloadManager, "_get_lora_scanner", AsyncMock(return_value=dummy_scanner)
)
monkeypatch.setattr(MetadataManager, "save_metadata", AsyncMock(return_value=True))
hash_calculator = AsyncMock(side_effect=["hash-one", "hash-two"])
monkeypatch.setattr(download_manager, "calculate_sha256", hash_calculator)
result = await manager._execute_download(
download_urls=download_urls,
@@ -372,15 +368,15 @@ async def test_execute_download_extracts_zip_multiple_models(monkeypatch, tmp_pa
assert extracted_one.exists()
assert extracted_two.exists()
assert hash_calculator.await_count == 2
assert MetadataManager.save_metadata.await_count == 2
assert dummy_scanner.add_model_to_cache.await_count == 2
metadata_calls = MetadataManager.save_metadata.await_args_list
assert metadata_calls[0].args[0] == str(extracted_one)
assert metadata_calls[0].args[1].sha256 == "hash-one"
# SHA256 comes from metadata (API value), not recalculated
assert metadata_calls[0].args[1].sha256 == "sha256"
assert metadata_calls[1].args[0] == str(extracted_two)
assert metadata_calls[1].args[1].sha256 == "hash-two"
assert metadata_calls[1].args[1].sha256 == "sha256"
@pytest.mark.asyncio
@@ -427,8 +423,6 @@ async def test_execute_download_extracts_zip_pt_embedding(monkeypatch, tmp_path)
ServiceRegistry, "get_embedding_scanner", AsyncMock(return_value=dummy_scanner)
)
monkeypatch.setattr(MetadataManager, "save_metadata", AsyncMock(return_value=True))
hash_calculator = AsyncMock(return_value="hash-pt")
monkeypatch.setattr(download_manager, "calculate_sha256", hash_calculator)
result = await manager._execute_download(
download_urls=download_urls,
@@ -445,10 +439,10 @@ async def test_execute_download_extracts_zip_pt_embedding(monkeypatch, tmp_path)
assert not zip_path.exists()
extracted = save_dir / "embedding.pt"
assert extracted.exists()
assert hash_calculator.await_args.args[0] == str(extracted)
saved_call = MetadataManager.save_metadata.await_args
assert saved_call.args[0] == str(extracted)
assert saved_call.args[1].sha256 == "hash-pt"
# SHA256 comes from metadata (API value), not recalculated
assert saved_call.args[1].sha256 == "sha256"
assert dummy_scanner.add_model_to_cache.await_count == 1

View File

@@ -9,7 +9,11 @@ from unittest.mock import AsyncMock, patch, MagicMock
import aiohttp
from py.services.downloader import Downloader, DownloadStalledError, DownloadRestartRequested
from py.services.downloader import (
Downloader,
DownloadStalledError,
DownloadRestartRequested,
)
class TestDownloadStreamControl:
@@ -118,6 +122,7 @@ class TestDownloaderConfiguration:
return instance1, instance2
import asyncio
instance1, instance2 = asyncio.run(get_instances())
assert instance1 is instance2
@@ -131,7 +136,7 @@ class TestDownloaderConfiguration:
downloader = Downloader()
assert downloader.chunk_size == 4 * 1024 * 1024 # 4MB
assert downloader.chunk_size == 16 * 1024 * 1024 # 16MB
assert downloader.max_retries == 5
assert downloader.base_delay == 2.0
assert downloader.session_timeout == 300
@@ -145,9 +150,9 @@ class TestDownloaderConfiguration:
downloader = Downloader()
assert 'User-Agent' in downloader.default_headers
assert 'ComfyUI-LoRA-Manager' in downloader.default_headers['User-Agent']
assert downloader.default_headers['Accept-Encoding'] == 'identity'
assert "User-Agent" in downloader.default_headers
assert "ComfyUI-LoRA-Manager" in downloader.default_headers["User-Agent"]
assert downloader.default_headers["Accept-Encoding"] == "identity"
# Cleanup
Downloader._instance = None
@@ -204,7 +209,10 @@ class TestDownloaderExceptions:
with pytest.raises(DownloadRestartRequested) as exc_info:
raise DownloadRestartRequested("Reconnect requested after resume")
assert "reconnect" in str(exc_info.value).lower() or "restart" in str(exc_info.value).lower()
assert (
"reconnect" in str(exc_info.value).lower()
or "restart" in str(exc_info.value).lower()
)
class TestDownloaderAuthHeaders:
@@ -217,8 +225,8 @@ class TestDownloaderAuthHeaders:
headers = downloader._get_auth_headers(use_auth=False)
assert 'User-Agent' in headers
assert 'Authorization' not in headers
assert "User-Agent" in headers
assert "Authorization" not in headers
Downloader._instance = None
@@ -231,12 +239,14 @@ class TestDownloaderAuthHeaders:
mock_settings = MagicMock()
mock_settings.get.return_value = None
with patch('py.services.downloader.get_settings_manager', return_value=mock_settings):
with patch(
"py.services.downloader.get_settings_manager", return_value=mock_settings
):
headers = downloader._get_auth_headers(use_auth=True)
# Should still have User-Agent but no Authorization
assert 'User-Agent' in headers
assert 'Authorization' not in headers
assert "User-Agent" in headers
assert "Authorization" not in headers
Downloader._instance = None
@@ -249,14 +259,16 @@ class TestDownloaderAuthHeaders:
mock_settings = MagicMock()
mock_settings.get.return_value = "test-api-key-12345"
with patch('py.services.downloader.get_settings_manager', return_value=mock_settings):
with patch(
"py.services.downloader.get_settings_manager", return_value=mock_settings
):
headers = downloader._get_auth_headers(use_auth=True)
# Should have both User-Agent and Authorization
assert 'User-Agent' in headers
assert 'Authorization' in headers
assert 'test-api-key-12345' in headers['Authorization']
assert headers['Content-Type'] == 'application/json'
assert "User-Agent" in headers
assert "Authorization" in headers
assert "test-api-key-12345" in headers["Authorization"]
assert headers["Content-Type"] == "application/json"
Downloader._instance = None
@@ -286,6 +298,7 @@ class TestDownloaderSessionManagement:
# Mock datetime to return current time
from datetime import datetime, timedelta
current_time = datetime.now()
downloader._session_created_at = current_time
@@ -301,6 +314,7 @@ class TestDownloaderSessionManagement:
# Simulate an old session (older than timeout)
from datetime import datetime, timedelta
old_time = datetime.now() - timedelta(seconds=downloader.session_timeout + 1)
downloader._session_created_at = old_time
downloader._session = MagicMock()