xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2026-11-10 18:53:31 +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 import aiohttp
from aiohttp import ClientResponse from aiohttp import ClientResponse
import itertools import itertools
import os import os
from unittest.mock import AsyncMock, patch, MagicMock 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 from model_filemanager import download_model, track_download_progress, create_model_path, check_file_exists, DownloadStatusType, DownloadModelStatus, validate_filename
import folder_paths
@pytest.fixture @pytest.fixture
def temp_dir(): def temp_dir():
@ -65,9 +66,12 @@ async def test_download_model_success(temp_dir):
mock_open.return_value.__enter__.return_value = mock_file mock_open.return_value.__enter__.return_value = mock_file
time_values = itertools.count(0, 0.1) 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')), \ 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('model_filemanager.check_file_exists', return_value=None), \
patch('builtins.open', mock_open), \ 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('folder_paths.get_folder_paths', return_value=[temp_dir]), \
patch('time.time', side_effect=time_values): # Simulate time passing 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_response.status = 404 # Simulate a "Not Found" error
mock_get = AsyncMock(return_value=mock_response) mock_get = AsyncMock(return_value=mock_response)
mock_progress_callback = AsyncMock() mock_progress_callback = AsyncMock()
fake_paths = {'checkpoints': (['mock_directory'], folder_paths.supported_pt_extensions)}
# Mock the create_model_path function # Mock the create_model_path function
with patch('model_filemanager.create_model_path', return_value='/mock/path/model.safetensors'), \ with patch('model_filemanager.create_model_path', return_value='/mock/path/model.safetensors'), \
patch('model_filemanager.check_file_exists', return_value=None), \ 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 # Call the function
result = await download_model( result = await download_model(
mock_get, mock_get,