Back to Repositories

Testing Internal Route Handlers and File Operations in ComfyUI

This test suite validates the internal routing functionality of ComfyUI’s server-side implementation, focusing on file system operations and route handling. The tests ensure proper API endpoint behavior for file listing operations across different directories.

Test Coverage Overview

The test suite provides comprehensive coverage of the internal routes implementation, particularly the file listing functionality.

Key areas tested include:

Valid directory file listing operations
Invalid directory handling
Error handling and edge cases
Route setup and application initialization
File service initialization with correct parameters

Implementation Analysis

The testing approach utilizes pytest’s async testing capabilities with aiohttp for HTTP request simulation. The implementation leverages fixture patterns for client setup and mock objects for isolating file system operations.

Technical patterns include:

Async/await testing patterns
Mock service injection
HTTP response validation
Application state verification

Technical Details

Testing tools and configuration:

pytest-asyncio for async test execution
aiohttp test client for API testing
unittest.mock for service mocking
pytest fixtures for test setup
Custom FileService implementation

Best Practices Demonstrated

The test suite exemplifies several testing best practices for Python web applications.

Notable practices include:

Isolation of external dependencies
Comprehensive error case coverage
Clear test organization and naming
Efficient fixture reuse
Proper async/await implementation

comfyanonymous/comfyui

tests-unit/server/routes/internal_routes_test.py

            
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(None)

@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(None)

        # 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