mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-03-26 23:48: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)
|
DownloadManager, "_get_lora_scanner", AsyncMock(return_value=dummy_scanner)
|
||||||
)
|
)
|
||||||
monkeypatch.setattr(MetadataManager, "save_metadata", AsyncMock(return_value=True))
|
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(
|
result = await manager._execute_download(
|
||||||
download_urls=download_urls,
|
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()
|
assert not zip_path.exists()
|
||||||
extracted = save_dir / "model.safetensors"
|
extracted = save_dir / "model.safetensors"
|
||||||
assert extracted.exists()
|
assert extracted.exists()
|
||||||
assert hash_calculator.await_args.args[0] == str(extracted)
|
|
||||||
saved_call = MetadataManager.save_metadata.await_args
|
saved_call = MetadataManager.save_metadata.await_args
|
||||||
assert saved_call.args[0] == str(extracted)
|
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
|
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)
|
DownloadManager, "_get_lora_scanner", AsyncMock(return_value=dummy_scanner)
|
||||||
)
|
)
|
||||||
monkeypatch.setattr(MetadataManager, "save_metadata", AsyncMock(return_value=True))
|
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(
|
result = await manager._execute_download(
|
||||||
download_urls=download_urls,
|
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_one.exists()
|
||||||
assert extracted_two.exists()
|
assert extracted_two.exists()
|
||||||
|
|
||||||
assert hash_calculator.await_count == 2
|
|
||||||
assert MetadataManager.save_metadata.await_count == 2
|
assert MetadataManager.save_metadata.await_count == 2
|
||||||
assert dummy_scanner.add_model_to_cache.await_count == 2
|
assert dummy_scanner.add_model_to_cache.await_count == 2
|
||||||
|
|
||||||
metadata_calls = MetadataManager.save_metadata.await_args_list
|
metadata_calls = MetadataManager.save_metadata.await_args_list
|
||||||
assert metadata_calls[0].args[0] == str(extracted_one)
|
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[0] == str(extracted_two)
|
||||||
assert metadata_calls[1].args[1].sha256 == "hash-two"
|
assert metadata_calls[1].args[1].sha256 == "sha256"
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@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)
|
ServiceRegistry, "get_embedding_scanner", AsyncMock(return_value=dummy_scanner)
|
||||||
)
|
)
|
||||||
monkeypatch.setattr(MetadataManager, "save_metadata", AsyncMock(return_value=True))
|
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(
|
result = await manager._execute_download(
|
||||||
download_urls=download_urls,
|
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()
|
assert not zip_path.exists()
|
||||||
extracted = save_dir / "embedding.pt"
|
extracted = save_dir / "embedding.pt"
|
||||||
assert extracted.exists()
|
assert extracted.exists()
|
||||||
assert hash_calculator.await_args.args[0] == str(extracted)
|
|
||||||
saved_call = MetadataManager.save_metadata.await_args
|
saved_call = MetadataManager.save_metadata.await_args
|
||||||
assert saved_call.args[0] == str(extracted)
|
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
|
assert dummy_scanner.add_model_to_cache.await_count == 1
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -9,7 +9,11 @@ from unittest.mock import AsyncMock, patch, MagicMock
|
|||||||
|
|
||||||
import aiohttp
|
import aiohttp
|
||||||
|
|
||||||
from py.services.downloader import Downloader, DownloadStalledError, DownloadRestartRequested
|
from py.services.downloader import (
|
||||||
|
Downloader,
|
||||||
|
DownloadStalledError,
|
||||||
|
DownloadRestartRequested,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
class TestDownloadStreamControl:
|
class TestDownloadStreamControl:
|
||||||
@@ -118,6 +122,7 @@ class TestDownloaderConfiguration:
|
|||||||
return instance1, instance2
|
return instance1, instance2
|
||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
|
|
||||||
instance1, instance2 = asyncio.run(get_instances())
|
instance1, instance2 = asyncio.run(get_instances())
|
||||||
|
|
||||||
assert instance1 is instance2
|
assert instance1 is instance2
|
||||||
@@ -131,7 +136,7 @@ class TestDownloaderConfiguration:
|
|||||||
|
|
||||||
downloader = Downloader()
|
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.max_retries == 5
|
||||||
assert downloader.base_delay == 2.0
|
assert downloader.base_delay == 2.0
|
||||||
assert downloader.session_timeout == 300
|
assert downloader.session_timeout == 300
|
||||||
@@ -145,9 +150,9 @@ class TestDownloaderConfiguration:
|
|||||||
|
|
||||||
downloader = Downloader()
|
downloader = Downloader()
|
||||||
|
|
||||||
assert 'User-Agent' in downloader.default_headers
|
assert "User-Agent" in downloader.default_headers
|
||||||
assert 'ComfyUI-LoRA-Manager' in downloader.default_headers['User-Agent']
|
assert "ComfyUI-LoRA-Manager" in downloader.default_headers["User-Agent"]
|
||||||
assert downloader.default_headers['Accept-Encoding'] == 'identity'
|
assert downloader.default_headers["Accept-Encoding"] == "identity"
|
||||||
|
|
||||||
# Cleanup
|
# Cleanup
|
||||||
Downloader._instance = None
|
Downloader._instance = None
|
||||||
@@ -204,7 +209,10 @@ class TestDownloaderExceptions:
|
|||||||
with pytest.raises(DownloadRestartRequested) as exc_info:
|
with pytest.raises(DownloadRestartRequested) as exc_info:
|
||||||
raise DownloadRestartRequested("Reconnect requested after resume")
|
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:
|
class TestDownloaderAuthHeaders:
|
||||||
@@ -217,8 +225,8 @@ class TestDownloaderAuthHeaders:
|
|||||||
|
|
||||||
headers = downloader._get_auth_headers(use_auth=False)
|
headers = downloader._get_auth_headers(use_auth=False)
|
||||||
|
|
||||||
assert 'User-Agent' in headers
|
assert "User-Agent" in headers
|
||||||
assert 'Authorization' not in headers
|
assert "Authorization" not in headers
|
||||||
|
|
||||||
Downloader._instance = None
|
Downloader._instance = None
|
||||||
|
|
||||||
@@ -231,12 +239,14 @@ class TestDownloaderAuthHeaders:
|
|||||||
mock_settings = MagicMock()
|
mock_settings = MagicMock()
|
||||||
mock_settings.get.return_value = None
|
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)
|
headers = downloader._get_auth_headers(use_auth=True)
|
||||||
|
|
||||||
# Should still have User-Agent but no Authorization
|
# Should still have User-Agent but no Authorization
|
||||||
assert 'User-Agent' in headers
|
assert "User-Agent" in headers
|
||||||
assert 'Authorization' not in headers
|
assert "Authorization" not in headers
|
||||||
|
|
||||||
Downloader._instance = None
|
Downloader._instance = None
|
||||||
|
|
||||||
@@ -249,14 +259,16 @@ class TestDownloaderAuthHeaders:
|
|||||||
mock_settings = MagicMock()
|
mock_settings = MagicMock()
|
||||||
mock_settings.get.return_value = "test-api-key-12345"
|
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)
|
headers = downloader._get_auth_headers(use_auth=True)
|
||||||
|
|
||||||
# Should have both User-Agent and Authorization
|
# Should have both User-Agent and Authorization
|
||||||
assert 'User-Agent' in headers
|
assert "User-Agent" in headers
|
||||||
assert 'Authorization' in headers
|
assert "Authorization" in headers
|
||||||
assert 'test-api-key-12345' in headers['Authorization']
|
assert "test-api-key-12345" in headers["Authorization"]
|
||||||
assert headers['Content-Type'] == 'application/json'
|
assert headers["Content-Type"] == "application/json"
|
||||||
|
|
||||||
Downloader._instance = None
|
Downloader._instance = None
|
||||||
|
|
||||||
@@ -286,6 +298,7 @@ class TestDownloaderSessionManagement:
|
|||||||
|
|
||||||
# Mock datetime to return current time
|
# Mock datetime to return current time
|
||||||
from datetime import datetime, timedelta
|
from datetime import datetime, timedelta
|
||||||
|
|
||||||
current_time = datetime.now()
|
current_time = datetime.now()
|
||||||
downloader._session_created_at = current_time
|
downloader._session_created_at = current_time
|
||||||
|
|
||||||
@@ -301,6 +314,7 @@ class TestDownloaderSessionManagement:
|
|||||||
|
|
||||||
# Simulate an old session (older than timeout)
|
# Simulate an old session (older than timeout)
|
||||||
from datetime import datetime, timedelta
|
from datetime import datetime, timedelta
|
||||||
|
|
||||||
old_time = datetime.now() - timedelta(seconds=downloader.session_timeout + 1)
|
old_time = datetime.now() - timedelta(seconds=downloader.session_timeout + 1)
|
||||||
downloader._session_created_at = old_time
|
downloader._session_created_at = old_time
|
||||||
downloader._session = MagicMock()
|
downloader._session = MagicMock()
|
||||||
|
|||||||
Reference in New Issue
Block a user