ComfyUI/tests-unit/utils/test_extra_custom_node_config.py

import pytest
import os
import sys
from unittest.mock import Mock, patch, mock_open
import yaml

from utils.extra_custom_node_config import load_extra_custom_node_path_config, get_current_custom_node_paths
import folder_paths

@pytest.fixture()
def clear_folder_paths():
    # Save original state
    original_paths = folder_paths.folder_names_and_paths.copy()
    original_custom_nodes = folder_paths.get_custom_nodes_directories()
    
    # Clear the custom nodes directories
    for path in original_custom_nodes:
        folder_paths.folder_names_and_paths["custom_nodes"][0].remove(path)
    
    yield
    
    # Restore original state
    folder_paths.folder_names_and_paths = original_paths

@pytest.fixture
def mock_yaml_content():
    return {
        'custom_nodes_config': {
            'base_path': '~/custom_nodes_dir',
        },
        'another_config': {
            'base_path': '/absolute/path/to/nodes'
        }
    }

@pytest.fixture
def yaml_config_with_vars():
    return """
    custom_nodes_config:
      base_path: '%APPDATA%/ComfyUI/custom_nodes'
    """

@patch('builtins.open', new_callable=mock_open, read_data="dummy file content")
@patch('os.path.expanduser')
@patch('yaml.safe_load')
def test_load_extra_custom_node_paths(
    mock_yaml_load, mock_expanduser, mock_file, mock_yaml_content, clear_folder_paths
):
    # Setup mocks
    mock_yaml_load.return_value = mock_yaml_content
    mock_expanduser.side_effect = lambda x: x.replace('~/', '/home/user/')
    
    # Mock add_custom_node_directory
    with patch('folder_paths.add_custom_node_directory') as mock_add_dir:
        load_extra_custom_node_path_config('dummy_path.yaml')
        
        # Verify the directories were added
        expected_paths = [
            '/home/user/custom_nodes_dir/custom_nodes',
            '/absolute/path/to/nodes/custom_nodes'
        ]
        
        # Check that add_custom_node_directory was called with the expected paths
        assert mock_add_dir.call_count == 2
        actual_paths = [call[0][0] for call in mock_add_dir.call_args_list]
        assert set(actual_paths) == set(expected_paths)

@patch('builtins.open', new_callable=mock_open, read_data="dummy file content")
@patch('os.path.expandvars')
@patch('yaml.safe_load')
def test_load_extra_custom_node_paths_with_env_vars(
    mock_yaml_load, mock_expandvars, mock_file, yaml_config_with_vars, clear_folder_paths
):
    # Setup mocks
    mock_yaml_load.return_value = yaml.safe_load(yaml_config_with_vars)
    
    def expandvars_side_effect(path):
        if '%APPDATA%' in path:
            if sys.platform == 'win32':
                return path.replace('%APPDATA%', 'C:\\Users\\TestUser\\AppData\\Roaming')
            else:
                return path.replace('%APPDATA%', '/Users/TestUser/AppData/Roaming')
        return path
    
    mock_expandvars.side_effect = expandvars_side_effect
    
    # Mock add_custom_node_directory
    with patch('folder_paths.add_custom_node_directory') as mock_add_dir:
        load_extra_custom_node_path_config('dummy_path.yaml')
        
        # Verify the directory was added with expanded path
        expected_path = os.path.join(
            expandvars_side_effect('%APPDATA%/ComfyUI/custom_nodes'),
            'custom_nodes'
        )
        mock_add_dir.assert_called_once_with(expected_path)

@patch('builtins.open', new_callable=mock_open, read_data="dummy file content")
@patch('logging.warning')
@patch('os.path.exists', return_value=False)
def test_load_extra_custom_node_paths_nonexistent(
    mock_exists, mock_warning, mock_file, mock_yaml_content, clear_folder_paths
):
    # Setup mocks
    with patch('yaml.safe_load', return_value=mock_yaml_content):
        with patch('os.path.expanduser', side_effect=lambda x: x.replace('~/', '/home/user/')):
            with patch('folder_paths.add_custom_node_directory') as mock_add_dir:
                load_extra_custom_node_path_config('dummy_path.yaml')
                
                # Verify warning was logged for non-existent paths
                assert mock_warning.call_count == 2
                for call in mock_warning.call_args_list:
                    assert "does not exist, skipping" in call[0][0]
                
                # Verify no directories were added
                mock_add_dir.assert_not_called()

def test_get_current_custom_node_paths(clear_folder_paths):
    # Add some test paths
    test_paths = ['/path/one', '/path/two']
    for path in test_paths:
        folder_paths.add_custom_node_directory(path)
    
    # Test getting the paths
    result = get_current_custom_node_paths()
    assert set(result) == set(test_paths)

@patch('builtins.open', side_effect=Exception("Test error"))
@patch('logging.error')
def test_load_extra_custom_node_paths_error(mock_error, mock_file, clear_folder_paths):
    # Test error handling when loading YAML fails
    load_extra_custom_node_path_config('invalid.yaml')
    mock_error.assert_called_once()
    assert "Failed to load extra custom node paths config from invalid.yaml" in str(mock_error.call_args[0][0])

@patch('builtins.open', new_callable=mock_open, read_data="dummy file content")
@patch('yaml.safe_load', return_value={})
@patch('logging.info')
def test_load_extra_custom_node_paths_empty_config(mock_info, mock_yaml_load, mock_file, clear_folder_paths):
    # Test with empty config
    load_extra_custom_node_path_config('empty.yaml')
    mock_info.assert_called_with("No custom node paths found in configuration")