sync with repo 28.08

tests-unit/app_test/frontend_manager_test.py

@@ -1,6 +1,7 @@
 import argparse
 import pytest
 from requests.exceptions import HTTPError
+from unittest.mock import patch
 
 from app.frontend_management import (
     FrontendManager,
@@ -83,6 +84,35 @@ def test_init_frontend_invalid_provider():
     with pytest.raises(HTTPError):
         FrontendManager.init_frontend_unsafe(version_string)
 
+@pytest.fixture
+def mock_os_functions():
+    with patch('app.frontend_management.os.makedirs') as mock_makedirs, \
+         patch('app.frontend_management.os.listdir') as mock_listdir, \
+         patch('app.frontend_management.os.rmdir') as mock_rmdir:
+        mock_listdir.return_value = []  # Simulate empty directory
+        yield mock_makedirs, mock_listdir, mock_rmdir
+
+@pytest.fixture
+def mock_download():
+    with patch('app.frontend_management.download_release_asset_zip') as mock:
+        mock.side_effect = Exception("Download failed")  # Simulate download failure
+        yield mock
+
+def test_finally_block(mock_os_functions, mock_download, mock_provider):
+    # Arrange
+    mock_makedirs, mock_listdir, mock_rmdir = mock_os_functions
+    version_string = 'test-owner/test-repo@1.0.0'
+
+    # Act & Assert
+    with pytest.raises(Exception):
+        FrontendManager.init_frontend_unsafe(version_string, mock_provider)
+
+    # Assert
+    mock_makedirs.assert_called_once()
+    mock_download.assert_called_once()
+    mock_listdir.assert_called_once()
+    mock_rmdir.assert_called_once()
+
 
 def test_parse_version_string():
     version_string = "owner/repo@1.0.0"

tests-unit/prompt_server_test/download_models_test.py

@@ -0,0 +1,321 @@
+import pytest
+import aiohttp
+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
+
+class AsyncIteratorMock:
+    """
+    A mock class that simulates an asynchronous iterator.
+    This is used to mimic the behavior of aiohttp's content iterator.
+    """
+    def __init__(self, seq):
+        # Convert the input sequence into an iterator
+        self.iter = iter(seq)
+
+    def __aiter__(self):
+        # This method is called when 'async for' is used
+        return self
+
+    async def __anext__(self):
+        # This method is called for each iteration in an 'async for' loop
+        try:
+            return next(self.iter)
+        except StopIteration:
+            # This is the asynchronous equivalent of StopIteration
+            raise StopAsyncIteration
+
+class ContentMock:
+    """
+    A mock class that simulates the content attribute of an aiohttp ClientResponse.
+    This class provides the iter_chunked method which returns an async iterator of chunks.
+    """
+    def __init__(self, chunks):
+        # Store the chunks that will be returned by the iterator
+        self.chunks = chunks
+
+    def iter_chunked(self, chunk_size):
+        # This method mimics aiohttp's content.iter_chunked()
+        # For simplicity in testing, we ignore chunk_size and just return our predefined chunks
+        return AsyncIteratorMock(self.chunks)
+
+@pytest.mark.asyncio
+async def test_download_model_success():
+    mock_response = AsyncMock(spec=aiohttp.ClientResponse)
+    mock_response.status = 200
+    mock_response.headers = {'Content-Length': '1000'}
+    # Create a mock for content that returns an async iterator directly
+    chunks = [b'a' * 500, b'b' * 300, b'c' * 200]
+    mock_response.content = ContentMock(chunks)
+
+    mock_make_request = AsyncMock(return_value=mock_response)
+    mock_progress_callback = AsyncMock()
+
+    # Mock file operations
+    mock_open = MagicMock()
+    mock_file = MagicMock()
+    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')), \
+         patch('model_filemanager.check_file_exists', return_value=None), \
+         patch('builtins.open', mock_open), \
+         patch('time.time', side_effect=time_values):  # Simulate time passing
+
+        result = await download_model(
+            mock_make_request,
+            'model.sft',
+            'http://example.com/model.sft',
+            'checkpoints',
+            mock_progress_callback
+        )
+
+    # Assert the result
+    assert isinstance(result, DownloadModelStatus)
+    assert result.message == 'Successfully downloaded model.sft'
+    assert result.status == 'completed'
+    assert result.already_existed is False
+
+    # Check progress callback calls
+    assert mock_progress_callback.call_count >= 3  # At least start, one progress update, and completion
+    
+    # Check initial call
+    mock_progress_callback.assert_any_call(
+        'checkpoints/model.sft',
+        DownloadModelStatus(DownloadStatusType.PENDING, 0, "Starting download of model.sft", False)
+    )
+
+    # Check final call
+    mock_progress_callback.assert_any_call(
+        'checkpoints/model.sft',
+        DownloadModelStatus(DownloadStatusType.COMPLETED, 100, "Successfully downloaded model.sft", False)
+    )
+
+    # Verify file writing
+    mock_file.write.assert_any_call(b'a' * 500)
+    mock_file.write.assert_any_call(b'b' * 300)
+    mock_file.write.assert_any_call(b'c' * 200)
+
+    # Verify request was made
+    mock_make_request.assert_called_once_with('http://example.com/model.sft')
+
+@pytest.mark.asyncio
+async def test_download_model_url_request_failure():
+    # Mock dependencies
+    mock_response = AsyncMock(spec=ClientResponse)
+    mock_response.status = 404  # Simulate a "Not Found" error
+    mock_get = AsyncMock(return_value=mock_response)
+    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')):
+        # 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
+            result = await download_model(
+                mock_get,
+                'model.safetensors',
+                'http://example.com/model.safetensors',
+                'mock_directory',
+                mock_progress_callback
+            )
+
+    # Assert the expected behavior
+    assert isinstance(result, DownloadModelStatus)
+    assert result.status == 'error'
+    assert result.message == 'Failed to download model.safetensors. Status code: 404'
+    assert result.already_existed is False
+
+    # Check that progress_callback was called with the correct arguments
+    mock_progress_callback.assert_any_call(
+        'mock_directory/model.safetensors',
+        DownloadModelStatus(
+            status=DownloadStatusType.PENDING,
+            progress_percentage=0,
+            message='Starting download of model.safetensors',
+            already_existed=False
+        )
+    )
+    mock_progress_callback.assert_called_with(
+        'mock_directory/model.safetensors',
+        DownloadModelStatus(
+            status=DownloadStatusType.ERROR,
+            progress_percentage=0,
+            message='Failed to download model.safetensors. Status code: 404',
+            already_existed=False
+        )
+    )
+
+    # Verify that the get method was called with the correct URL
+    mock_get.assert_called_once_with('http://example.com/model.safetensors')
+
+@pytest.mark.asyncio
+async def test_download_model_invalid_model_subdirectory():
+    
+    mock_make_request = AsyncMock()
+    mock_progress_callback = AsyncMock()
+
+    
+    result = await download_model(
+        mock_make_request,
+        'model.sft',
+        'http://example.com/model.sft',
+        '../bad_path',
+        mock_progress_callback
+    )
+
+    # Assert the result
+    assert isinstance(result, DownloadModelStatus)
+    assert result.message == 'Invalid model subdirectory'
+    assert result.status == 'error'
+    assert result.already_existed is False
+
+
+# For create_model_path function
+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)
+    
+    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))
+
+
+@pytest.mark.asyncio
+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")
+    
+    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",
+        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")
+    
+    assert result is None
+    mock_callback.assert_not_called()
+
+@pytest.mark.asyncio
+async def test_track_download_progress_no_content_length():
+    mock_response = AsyncMock(spec=aiohttp.ClientResponse)
+    mock_response.headers = {}  # No Content-Length header
+    mock_response.content.iter_chunked.return_value = AsyncIteratorMock([b'a' * 500, b'b' * 500])
+
+    mock_callback = AsyncMock()
+    mock_open = MagicMock(return_value=MagicMock())
+
+    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
+        )
+
+    assert result.status == "completed"
+    # Check that progress was reported even without knowing the total size
+    mock_callback.assert_any_call(
+        'models/model.sft',
+        DownloadModelStatus(DownloadStatusType.IN_PROGRESS, 0, "Downloading model.sft", already_existed=False)
+    )
+
+@pytest.mark.asyncio
+async def test_track_download_progress_interval():
+    mock_response = AsyncMock(spec=aiohttp.ClientResponse)
+    mock_response.headers = {'Content-Length': '1000'}
+    mock_response.content.iter_chunked.return_value = AsyncIteratorMock([b'a' * 100] * 10)
+
+    mock_callback = AsyncMock()
+    mock_open = MagicMock(return_value=MagicMock())
+
+    # Create a mock time function that returns incremental float values
+    mock_time = MagicMock()
+    mock_time.side_effect = [i * 0.5 for i in range(30)]  # This should be enough for 10 chunks
+
+    with patch('builtins.open', mock_open), \
+         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
+        )
+
+    # Print out the actual call count and the arguments of each call for debugging
+    print(f"mock_callback was called {mock_callback.call_count} times")
+    for i, call in enumerate(mock_callback.call_args_list):
+        args, kwargs = call
+        print(f"Call {i + 1}: {args[1].status}, Progress: {args[1].progress_percentage:.2f}%")
+
+    # Assert that progress was updated at least 3 times (start, at least one interval, and end)
+    assert mock_callback.call_count >= 3, f"Expected at least 3 calls, but got {mock_callback.call_count}"
+
+    # Verify the first and last calls
+    first_call = mock_callback.call_args_list[0]
+    assert first_call[0][1].status == "in_progress"
+    # Allow for some initial progress, but it should be less than 50%
+    assert 0 <= first_call[0][1].progress_percentage < 50, f"First call progress was {first_call[0][1].progress_percentage}%"
+
+    last_call = mock_callback.call_args_list[-1]
+    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),
+    ("valid model.safetensors", True), # Test with space
+    ("UPPERCASE_MODEL.SAFETENSORS", True),
+    ("model_with.multiple.dots.pt", False),
+    ("", False),  # Empty string
+    ("../../../etc/passwd", False),  # Path traversal attempt
+    ("/etc/passwd", False),  # Absolute path
+    ("\\windows\\system32\\config\\sam", False),  # Windows path
+    (".hidden_file.pt", False),  # Hidden file
+    ("invalid<char>.ckpt", False),  # Invalid character
+    ("invalid?.ckpt", False),  # Another invalid character
+    ("very" * 100 + ".safetensors", False),  # Too long filename
+    ("\nmodel_with_newline.pt", False),  # Newline character
+    ("model_with_emoji😊.pt", False),  # Emoji in filename
+])
+def test_validate_filename(filename, expected):
+    assert validate_filename(filename) == expected

tests-unit/requirements.txt

@@ -1 +1,3 @@
 pytest>=7.8.0
+pytest-aiohttp
+pytest-asyncio

tests-unit/server/routes/internal_routes_test.py

@@ -0,0 +1,115 @@
+import pytest
+from aiohttp import web
+from unittest.mock import MagicMock, patch
+from api_server.routes.internal.internal_routes import InternalRoutes
+from api_server.services.file_service import FileService
+from folder_paths import models_dir, user_directory, output_directory
+
+
+@pytest.fixture
+def internal_routes():
+    return InternalRoutes()
+
+@pytest.fixture
+def aiohttp_client_factory(aiohttp_client, internal_routes):
+    async def _get_client():
+        app = internal_routes.get_app()
+        return await aiohttp_client(app)
+    return _get_client
+
+@pytest.mark.asyncio
+async def test_list_files_valid_directory(aiohttp_client_factory, internal_routes):
+    mock_file_list = [
+        {"name": "file1.txt", "path": "file1.txt", "type": "file", "size": 100},
+        {"name": "dir1", "path": "dir1", "type": "directory"}
+    ]
+    internal_routes.file_service.list_files = MagicMock(return_value=mock_file_list)
+    client = await aiohttp_client_factory()
+    resp = await client.get('/files?directory=models')
+    assert resp.status == 200
+    data = await resp.json()
+    assert 'files' in data
+    assert len(data['files']) == 2
+    assert data['files'] == mock_file_list
+
+    # Check other valid directories
+    resp = await client.get('/files?directory=user')
+    assert resp.status == 200
+    resp = await client.get('/files?directory=output')
+    assert resp.status == 200
+
+@pytest.mark.asyncio
+async def test_list_files_invalid_directory(aiohttp_client_factory, internal_routes):
+    internal_routes.file_service.list_files = MagicMock(side_effect=ValueError("Invalid directory key"))
+    client = await aiohttp_client_factory()
+    resp = await client.get('/files?directory=invalid')
+    assert resp.status == 400
+    data = await resp.json()
+    assert 'error' in data
+    assert data['error'] == "Invalid directory key"
+
+@pytest.mark.asyncio
+async def test_list_files_exception(aiohttp_client_factory, internal_routes):
+    internal_routes.file_service.list_files = MagicMock(side_effect=Exception("Unexpected error"))
+    client = await aiohttp_client_factory()
+    resp = await client.get('/files?directory=models')
+    assert resp.status == 500
+    data = await resp.json()
+    assert 'error' in data
+    assert data['error'] == "Unexpected error"
+
+@pytest.mark.asyncio
+async def test_list_files_no_directory_param(aiohttp_client_factory, internal_routes):
+    mock_file_list = []
+    internal_routes.file_service.list_files = MagicMock(return_value=mock_file_list)
+    client = await aiohttp_client_factory()
+    resp = await client.get('/files')
+    assert resp.status == 200
+    data = await resp.json()
+    assert 'files' in data
+    assert len(data['files']) == 0
+
+def test_setup_routes(internal_routes):
+    internal_routes.setup_routes()
+    routes = internal_routes.routes
+    assert any(route.method == 'GET' and str(route.path) == '/files' for route in routes)
+
+def test_get_app(internal_routes):
+    app = internal_routes.get_app()
+    assert isinstance(app, web.Application)
+    assert internal_routes._app is not None
+
+def test_get_app_reuse(internal_routes):
+    app1 = internal_routes.get_app()
+    app2 = internal_routes.get_app()
+    assert app1 is app2
+
+@pytest.mark.asyncio
+async def test_routes_added_to_app(aiohttp_client_factory, internal_routes):
+    client = await aiohttp_client_factory()
+    try:
+        resp = await client.get('/files')
+        print(f"Response received: status {resp.status}")
+    except Exception as e:
+        print(f"Exception occurred during GET request: {e}")
+        raise
+
+    assert resp.status != 404, "Route /files does not exist"
+
+@pytest.mark.asyncio
+async def test_file_service_initialization():
+    with patch('api_server.routes.internal.internal_routes.FileService') as MockFileService:
+        # Create a mock instance
+        mock_file_service_instance = MagicMock(spec=FileService)
+        MockFileService.return_value = mock_file_service_instance
+        internal_routes = InternalRoutes()
+
+        # Check if FileService was initialized with the correct parameters
+        MockFileService.assert_called_once_with({
+            "models": models_dir,
+            "user": user_directory,
+            "output": output_directory
+        })
+
+        # Verify that the file_service attribute of InternalRoutes is set
+        assert internal_routes.file_service == mock_file_service_instance

tests-unit/server/services/file_service_test.py

@@ -0,0 +1,54 @@
+import pytest
+from unittest.mock import MagicMock
+from api_server.services.file_service import FileService
+
+@pytest.fixture
+def mock_file_system_ops():
+    return MagicMock()
+
+@pytest.fixture
+def file_service(mock_file_system_ops):
+    allowed_directories = {
+        "models": "/path/to/models",
+        "user": "/path/to/user",
+        "output": "/path/to/output"
+    }
+    return FileService(allowed_directories, file_system_ops=mock_file_system_ops)
+
+def test_list_files_valid_directory(file_service, mock_file_system_ops):
+    mock_file_system_ops.walk_directory.return_value = [
+        {"name": "file1.txt", "path": "file1.txt", "type": "file", "size": 100},
+        {"name": "dir1", "path": "dir1", "type": "directory"}
+    ]
+    
+    result = file_service.list_files("models")
+    
+    assert len(result) == 2
+    assert result[0]["name"] == "file1.txt"
+    assert result[1]["name"] == "dir1"
+    mock_file_system_ops.walk_directory.assert_called_once_with("/path/to/models")
+
+def test_list_files_invalid_directory(file_service):
+    # Does not support walking directories outside of the allowed directories
+    with pytest.raises(ValueError, match="Invalid directory key"):
+        file_service.list_files("invalid_key")
+
+def test_list_files_empty_directory(file_service, mock_file_system_ops):
+    mock_file_system_ops.walk_directory.return_value = []
+    
+    result = file_service.list_files("models")
+    
+    assert len(result) == 0
+    mock_file_system_ops.walk_directory.assert_called_once_with("/path/to/models")
+
+@pytest.mark.parametrize("directory_key", ["models", "user", "output"])
+def test_list_files_all_allowed_directories(file_service, mock_file_system_ops, directory_key):
+    mock_file_system_ops.walk_directory.return_value = [
+        {"name": f"file_{directory_key}.txt", "path": f"file_{directory_key}.txt", "type": "file", "size": 100}
+    ]
+    
+    result = file_service.list_files(directory_key)
+    
+    assert len(result) == 1
+    assert result[0]["name"] == f"file_{directory_key}.txt"
+    mock_file_system_ops.walk_directory.assert_called_once_with(f"/path/to/{directory_key}")

tests-unit/server/utils/file_operations_test.py

@@ -0,0 +1,42 @@
+import pytest
+from typing import List
+from api_server.utils.file_operations import FileSystemOperations, FileSystemItem, is_file_info
+
+@pytest.fixture
+def temp_directory(tmp_path):
+    # Create a temporary directory structure
+    dir1 = tmp_path / "dir1"
+    dir2 = tmp_path / "dir2"
+    dir1.mkdir()
+    dir2.mkdir()
+    (dir1 / "file1.txt").write_text("content1")
+    (dir2 / "file2.txt").write_text("content2")
+    (tmp_path / "file3.txt").write_text("content3")
+    return tmp_path
+
+def test_walk_directory(temp_directory):
+    result: List[FileSystemItem] = FileSystemOperations.walk_directory(str(temp_directory))
+    
+    assert len(result) == 5  # 2 directories and 3 files
+    
+    files = [item for item in result if item['type'] == 'file']
+    dirs = [item for item in result if item['type'] == 'directory']
+    
+    assert len(files) == 3
+    assert len(dirs) == 2
+    
+    file_names = {file['name'] for file in files}
+    assert file_names == {'file1.txt', 'file2.txt', 'file3.txt'}
+    
+    dir_names = {dir['name'] for dir in dirs}
+    assert dir_names == {'dir1', 'dir2'}
+
+def test_walk_directory_empty(tmp_path):
+    result = FileSystemOperations.walk_directory(str(tmp_path))
+    assert len(result) == 0
+
+def test_walk_directory_file_size(temp_directory):
+    result: List[FileSystemItem] = FileSystemOperations.walk_directory(str(temp_directory))
+    files = [item for item in result if is_file_info(item)]
+    for file in files:
+        assert file['size'] > 0  # Assuming all files have some content