mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-03-26 15:38:52 -03:00
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:
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user