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

test tweaks

Browse Source
This commit is contained in:
Alex "mcmonkey" Goodwin 2024-09-19 20:45:27 +09:00
parent 3e8dbd51da
commit 57a91a2c1c
2 changed files with 10 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
model_filemanager/download_models.py
View File

@ -3,7 +3,7 @@ import aiohttp
import os import os
import traceback import traceback
import logging import logging
from folder_paths import models_dir, folder_names_and_paths, get_folder_paths from folder_paths import folder_names_and_paths, get_folder_paths
import re import re
from typing import Callable, Any, Optional, Awaitable, Dict from typing import Callable, Any, Optional, Awaitable, Dict
from enum import Enum from enum import Enum

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

@ -1,4 +1,5 @@
import pytest import pytest
import tempfile
import aiohttp import aiohttp
from aiohttp import ClientResponse from aiohttp import ClientResponse
import itertools import itertools
@ -6,6 +7,11 @@ 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
@pytest.fixture
def temp_dir():
with tempfile.TemporaryDirectory() as tmpdirname:
yield tmpdirname
class AsyncIteratorMock: class AsyncIteratorMock:
""" """
A mock class that simulates an asynchronous iterator. A mock class that simulates an asynchronous iterator.
@ -42,7 +48,7 @@ class ContentMock:
return AsyncIteratorMock(self.chunks) return AsyncIteratorMock(self.chunks)
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_download_model_success(): async def test_download_model_success(temp_dir):
mock_response = AsyncMock(spec=aiohttp.ClientResponse) mock_response = AsyncMock(spec=aiohttp.ClientResponse)
mock_response.status = 200 mock_response.status = 200
mock_response.headers = {'Content-Length': '1000'} mock_response.headers = {'Content-Length': '1000'}
@ -62,7 +68,7 @@ async def test_download_model_success():
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.get_folder_paths', return_value=['mock_directory']), \ 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
result = await download_model( result = await download_model(
@ -70,7 +76,7 @@ async def test_download_model_success():
'model.sft', 'model.sft',
'http://example.com/model.sft', 'http://example.com/model.sft',
'checkpoints', 'checkpoints',
'mock_directory', temp_dir,
mock_progress_callback mock_progress_callback
) )