wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-07-21 23:41:28 +08:00
Code Issues Actions 10 Packages Projects Releases Wiki Activity

Fix tests

Browse Source
This commit is contained in:
doctorpangloss
2024-10-29 14:27:14 -07:00
parent 4a13766d14
commit b42e59d602
3 changed files with 105 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/unit/prompt_server_test/download_models_test.py
+6 -5
View File
@@ -9,6 +9,7 @@ from aiohttp import ClientResponse
from comfy.cmd import folder_paths
from comfy.component_model.folder_path_types import FolderNames, FolderPathsTuple, supported_pt_extensions
from comfy.execution_context import context_folder_names_and_paths
from comfy.model_filemanager import download_model, track_download_progress, \
create_model_path, check_file_exists, DownloadStatusType, DownloadModelStatus, validate_filename
@@ -74,9 +75,9 @@ async def test_download_model_success(temp_dir):
fake_paths = FolderNames.from_dict({'checkpoints': ([temp_dir], folder_paths.supported_pt_extensions)})
with patch('comfy.model_filemanager.create_model_path', return_value=('models/checkpoints/model.sft', 'model.sft')), \
with context_folder_names_and_paths(fake_paths), \
patch('comfy.model_filemanager.create_model_path', return_value=('models/checkpoints/model.sft', 'model.sft')), \
patch('comfy.model_filemanager.check_file_exists', return_value=None), \
patch('comfy.cmd.folder_paths.folder_names_and_paths', fake_paths), \
patch('time.time', side_effect=time_values): # Simulate time passing
result = await download_model(
@@ -130,9 +131,9 @@ async def test_download_model_url_request_failure(temp_dir):
fake_paths = FolderNames.from_dict({'checkpoints': ([temp_dir], folder_paths.supported_pt_extensions)})
# Mock the create_model_path function
with patch('comfy.model_filemanager.create_model_path', return_value='/mock/path/model.safetensors'), \
patch('comfy.model_filemanager.check_file_exists', return_value=None), \
patch('comfy.cmd.folder_paths.folder_names_and_paths', fake_paths):
with context_folder_names_and_paths(fake_paths), \
patch('comfy.model_filemanager.create_model_path', return_value='/mock/path/model.safetensors'), \
patch('comfy.model_filemanager.check_file_exists', return_value=None):
# Call the function
result = await download_model(
mock_get,