diff --git a/tests-unit/prompt_server_test/download_models_test.py b/tests-unit/prompt_server_test/download_models_test.py index 66150a468..e86ba2d76 100644 --- a/tests-unit/prompt_server_test/download_models_test.py +++ b/tests-unit/prompt_server_test/download_models_test.py @@ -4,7 +4,7 @@ from aiohttp import ClientResponse import itertools import os from unittest.mock import AsyncMock, patch, MagicMock -from model_filemanager import download_model, validate_model_subdirectory, 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 class AsyncIteratorMock: """ @@ -59,7 +59,7 @@ async def test_download_model_success(): mock_open.return_value.__enter__.return_value = mock_file time_values = itertools.count(0, 0.1) - with patch('model_filemanager.create_model_path', return_value=('models/checkpoints/model.sft', 'checkpoints/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('builtins.open', mock_open), \ patch('time.time', side_effect=time_values): # Simulate time passing @@ -69,6 +69,7 @@ async def test_download_model_success(): 'model.sft', 'http://example.com/model.sft', 'checkpoints', + 'mock_directory', mock_progress_callback ) @@ -83,13 +84,13 @@ async def test_download_model_success(): # Check initial call mock_progress_callback.assert_any_call( - 'checkpoints/model.sft', + 'model.sft', DownloadModelStatus(DownloadStatusType.PENDING, 0, "Starting download of model.sft", False) ) # Check final call mock_progress_callback.assert_any_call( - 'checkpoints/model.sft', + 'model.sft', DownloadModelStatus(DownloadStatusType.COMPLETED, 100, "Successfully downloaded model.sft", False) ) @@ -110,7 +111,7 @@ async def test_download_model_url_request_failure(): mock_progress_callback = AsyncMock() # Mock the create_model_path function - with patch('model_filemanager.create_model_path', return_value=('/mock/path/model.safetensors', 'mock/path/model.safetensors')): + with patch('model_filemanager.create_model_path', return_value='/mock/path/model.safetensors'): # Mock the check_file_exists function to return None (file doesn't exist) with patch('model_filemanager.check_file_exists', return_value=None): # Call the function @@ -118,6 +119,7 @@ async def test_download_model_url_request_failure(): mock_get, 'model.safetensors', 'http://example.com/model.safetensors', + 'checkpoints', 'mock_directory', mock_progress_callback ) @@ -163,12 +165,13 @@ async def test_download_model_invalid_model_subdirectory(): 'model.sft', 'http://example.com/model.sft', '../bad_path', + '../bad_path', mock_progress_callback ) # Assert the result assert isinstance(result, DownloadModelStatus) - assert result.message == 'Invalid model subdirectory' + assert result.message.startswith('Invalid or unrecognized model directory') assert result.status == 'error' assert result.already_existed is False @@ -177,14 +180,13 @@ async def test_download_model_invalid_model_subdirectory(): def test_create_model_path(tmp_path, monkeypatch): mock_models_dir = tmp_path / "models" monkeypatch.setattr('folder_paths.models_dir', str(mock_models_dir)) - + model_name = "test_model.sft" model_directory = "test_dir" - - file_path, relative_path = create_model_path(model_name, model_directory, mock_models_dir) - + + file_path = create_model_path(model_name, model_directory, mock_models_dir) + assert file_path == str(mock_models_dir / model_directory / model_name) - assert relative_path == f"{model_directory}/{model_name}" assert os.path.exists(os.path.dirname(file_path)) @@ -192,29 +194,29 @@ def test_create_model_path(tmp_path, monkeypatch): async def test_check_file_exists_when_file_exists(tmp_path): file_path = tmp_path / "existing_model.sft" file_path.touch() # Create an empty file - + mock_callback = AsyncMock() - - result = await check_file_exists(str(file_path), "existing_model.sft", mock_callback, "test/existing_model.sft") - + + result = await check_file_exists(str(file_path), "existing_model.sft", mock_callback) + assert result is not None assert result.status == "completed" assert result.message == "existing_model.sft already exists" assert result.already_existed is True - + mock_callback.assert_called_once_with( - "test/existing_model.sft", + "existing_model.sft", DownloadModelStatus(DownloadStatusType.COMPLETED, 100, "existing_model.sft already exists", already_existed=True) ) @pytest.mark.asyncio async def test_check_file_exists_when_file_does_not_exist(tmp_path): file_path = tmp_path / "non_existing_model.sft" - + mock_callback = AsyncMock() - - result = await check_file_exists(str(file_path), "non_existing_model.sft", mock_callback, "test/non_existing_model.sft") - + + result = await check_file_exists(str(file_path), "non_existing_model.sft", mock_callback) + assert result is None mock_callback.assert_not_called() @@ -230,13 +232,13 @@ async def test_track_download_progress_no_content_length(): with patch('builtins.open', mock_open): result = await track_download_progress( mock_response, '/mock/path/model.sft', 'model.sft', - mock_callback, 'models/model.sft', interval=0.1 + mock_callback, interval=0.1 ) assert result.status == "completed" # Check that progress was reported even without knowing the total size mock_callback.assert_any_call( - 'models/model.sft', + 'model.sft', DownloadModelStatus(DownloadStatusType.IN_PROGRESS, 0, "Downloading model.sft", already_existed=False) ) @@ -257,7 +259,7 @@ async def test_track_download_progress_interval(): patch('time.time', mock_time): await track_download_progress( mock_response, '/mock/path/model.sft', 'model.sft', - mock_callback, 'models/model.sft', interval=1.0 + mock_callback, interval=1.0 ) # Print out the actual call count and the arguments of each call for debugging @@ -279,27 +281,6 @@ async def test_track_download_progress_interval(): assert last_call[0][1].status == "completed" assert last_call[0][1].progress_percentage == 100 -def test_valid_subdirectory(): - assert validate_model_subdirectory("valid-model123") is True - -def test_subdirectory_too_long(): - assert validate_model_subdirectory("a" * 51) is False - -def test_subdirectory_with_double_dots(): - assert validate_model_subdirectory("model/../unsafe") is False - -def test_subdirectory_with_slash(): - assert validate_model_subdirectory("model/unsafe") is False - -def test_subdirectory_with_special_characters(): - assert validate_model_subdirectory("model@unsafe") is False - -def test_subdirectory_with_underscore_and_dash(): - assert validate_model_subdirectory("valid_model-name") is True - -def test_empty_subdirectory(): - assert validate_model_subdirectory("") is False - @pytest.mark.parametrize("filename, expected", [ ("valid_model.safetensors", True), ("valid_model.sft", True),