partial tests update

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),