xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2026-11-09 10:26:38 +08:00
Code Issues Packages Projects Releases Wiki Activity

workaround the first tests blocker

Browse Source
This commit is contained in:
Alex "mcmonkey" Goodwin 2024-09-22 16:38:38 +09:00
parent db1814cc44
commit 3ac444c35f
1 changed files with 9 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
tests-unit/prompt_server_test/download_models_test.py
View File

@ -3,9 +3,10 @@ import tempfile
import aiohttp
from aiohttp import ClientResponse
import itertools
import os
import os
from unittest.mock import AsyncMock, patch, MagicMock
from model_filemanager import download_model, track_download_progress, create_model_path, check_file_exists, DownloadStatusType, DownloadModelStatus, validate_filename
import folder_paths
@pytest.fixture
def temp_dir():
@ -65,9 +66,12 @@ async def test_download_model_success(temp_dir):
mock_open.return_value.__enter__.return_value = mock_file
time_values = itertools.count(0, 0.1)
fake_paths = {'checkpoints': ([temp_dir], folder_paths.supported_pt_extensions)}
with patch('model_filemanager.create_model_path', return_value=('models/checkpoints/model.sft', 'model.sft')), \
patch('model_filemanager.check_file_exists', return_value=None), \
patch('builtins.open', mock_open), \
patch('folder_paths.folder_names_and_paths', fake_paths), \
patch('folder_paths.get_folder_paths', return_value=[temp_dir]), \
patch('time.time', side_effect=time_values): # Simulate time passing
@ -116,11 +120,14 @@ async def test_download_model_url_request_failure():
mock_response.status = 404 # Simulate a "Not Found" error
mock_get = AsyncMock(return_value=mock_response)
mock_progress_callback = AsyncMock()
fake_paths = {'checkpoints': (['mock_directory'], folder_paths.supported_pt_extensions)}
# Mock the create_model_path function
with patch('model_filemanager.create_model_path', return_value='/mock/path/model.safetensors'), \
patch('model_filemanager.check_file_exists', return_value=None), \
patch('folder_paths.get_folder_paths', return_value=['mock_directory']):
patch('folder_paths.get_folder_paths', return_value=['mock_directory']), \
patch('folder_paths.folder_names_and_paths', fake_paths):
# Call the function
result = await download_model(
mock_get,