avoid 'mock_directory'

use temp dir to avoid accidental fs pollution from tests
This commit is contained in:
Alex "mcmonkey" Goodwin 2024-09-23 19:00:18 +09:00
parent ae931da097
commit 89b3539dd6

View File

@ -109,19 +109,18 @@ async def test_download_model_success(temp_dir):
mock_make_request.assert_called_once_with('http://example.com/model.sft') mock_make_request.assert_called_once_with('http://example.com/model.sft')
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_download_model_url_request_failure(): async def test_download_model_url_request_failure(temp_dir):
# Mock dependencies # Mock dependencies
mock_response = AsyncMock(spec=ClientResponse) mock_response = AsyncMock(spec=ClientResponse)
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)} fake_paths = {'checkpoints': ([temp_dir], 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.folder_names_and_paths', fake_paths): patch('folder_paths.folder_names_and_paths', fake_paths):
# Call the function # Call the function
result = await download_model( result = await download_model(
@ -129,7 +128,7 @@ async def test_download_model_url_request_failure():
'model.safetensors', 'model.safetensors',
'http://example.com/model.safetensors', 'http://example.com/model.safetensors',
'checkpoints', 'checkpoints',
'mock_directory', temp_dir,
mock_progress_callback mock_progress_callback
) )