fred-dev/comfy_ui_ali

1import pytest2import yaml3import os4import sys5from unittest.mock import Mock, patch, mock_open6 7from utils.extra_config import load_extra_path_config8import folder_paths9 10 11@pytest.fixture()12def clear_folder_paths():13    # Clear the global dictionary before each test to ensure isolation14    original = folder_paths.folder_names_and_paths.copy()15    folder_paths.folder_names_and_paths.clear()16    yield17    folder_paths.folder_names_and_paths = original18 19 20@pytest.fixture21def mock_yaml_content():22    return {23        'test_config': {24            'base_path': '~/App/',25            'checkpoints': 'subfolder1',26        }27    }28 29 30@pytest.fixture31def mock_expanded_home():32    return '/home/user'33 34 35@pytest.fixture36def yaml_config_with_appdata():37    return """38    test_config:39      base_path: '%APPDATA%/ComfyUI'40      checkpoints: 'models/checkpoints'41    """42 43 44@pytest.fixture45def mock_yaml_content_appdata(yaml_config_with_appdata):46    return yaml.safe_load(yaml_config_with_appdata)47 48 49@pytest.fixture50def mock_expandvars_appdata():51    mock = Mock()52 53    def expandvars(path):54        if '%APPDATA%' in path:55            if sys.platform == 'win32':56                return path.replace('%APPDATA%', 'C:/Users/TestUser/AppData/Roaming')57            else:58                return path.replace('%APPDATA%', '/Users/TestUser/AppData/Roaming')59        return path60 61    mock.side_effect = expandvars62    return mock63 64 65@pytest.fixture66def mock_add_model_folder_path():67    return Mock()68 69 70@pytest.fixture71def mock_expanduser(mock_expanded_home):72    def _expanduser(path):73        if path.startswith('~/'):74            return os.path.join(mock_expanded_home, path[2:])75        return path76    return _expanduser77 78 79@pytest.fixture80def mock_yaml_safe_load(mock_yaml_content):81    return Mock(return_value=mock_yaml_content)82 83 84@patch('builtins.open', new_callable=mock_open, read_data="dummy file content")85def test_load_extra_model_paths_expands_userpath(86    mock_file,87    monkeypatch,88    mock_add_model_folder_path,89    mock_expanduser,90    mock_yaml_safe_load,91    mock_expanded_home92):93    # Attach mocks used by load_extra_path_config94    monkeypatch.setattr(folder_paths, 'add_model_folder_path', mock_add_model_folder_path)95    monkeypatch.setattr(os.path, 'expanduser', mock_expanduser)96    monkeypatch.setattr(yaml, 'safe_load', mock_yaml_safe_load)97 98    dummy_yaml_file_name = 'dummy_path.yaml'99    load_extra_path_config(dummy_yaml_file_name)100 101    expected_calls = [102        ('checkpoints', os.path.join(mock_expanded_home, 'App', 'subfolder1'), False),103    ]104 105    assert mock_add_model_folder_path.call_count == len(expected_calls)106 107    # Check if add_model_folder_path was called with the correct arguments108    for actual_call, expected_call in zip(mock_add_model_folder_path.call_args_list, expected_calls):109        assert actual_call.args[0] == expected_call[0]110        assert os.path.normpath(actual_call.args[1]) == os.path.normpath(expected_call[1])  # Normalize and check the path to check on multiple OS.111        assert actual_call.args[2] == expected_call[2]112 113    # Check if yaml.safe_load was called114    mock_yaml_safe_load.assert_called_once()115 116    # Check if open was called with the correct file path117    mock_file.assert_called_once_with(dummy_yaml_file_name, 'r', encoding='utf-8')118 119 120@patch('builtins.open', new_callable=mock_open)121def test_load_extra_model_paths_expands_appdata(122    mock_file,123    monkeypatch,124    mock_add_model_folder_path,125    mock_expandvars_appdata,126    yaml_config_with_appdata,127    mock_yaml_content_appdata128):129    # Set the mock_file to return yaml with appdata as a variable130    mock_file.return_value.read.return_value = yaml_config_with_appdata131 132    # Attach mocks133    monkeypatch.setattr(folder_paths, 'add_model_folder_path', mock_add_model_folder_path)134    monkeypatch.setattr(os.path, 'expandvars', mock_expandvars_appdata)135    monkeypatch.setattr(yaml, 'safe_load', Mock(return_value=mock_yaml_content_appdata))136 137    # Mock expanduser to do nothing (since we're not testing it here)138    monkeypatch.setattr(os.path, 'expanduser', lambda x: x)139 140    dummy_yaml_file_name = 'dummy_path.yaml'141    load_extra_path_config(dummy_yaml_file_name)142 143    if sys.platform == "win32":144        expected_base_path = 'C:/Users/TestUser/AppData/Roaming/ComfyUI'145    else:146        expected_base_path = '/Users/TestUser/AppData/Roaming/ComfyUI'147    expected_calls = [148        ('checkpoints', os.path.normpath(os.path.join(expected_base_path, 'models/checkpoints')), False),149    ]150 151    assert mock_add_model_folder_path.call_count == len(expected_calls)152 153    # Check the base path variable was expanded154    for actual_call, expected_call in zip(mock_add_model_folder_path.call_args_list, expected_calls):155        assert actual_call.args == expected_call156 157    # Verify that expandvars was called158    assert mock_expandvars_appdata.called159 160 161@patch("builtins.open", new_callable=mock_open, read_data="dummy yaml content")162@patch("yaml.safe_load")163def test_load_extra_path_config_relative_base_path(164    mock_yaml_load, _mock_file, clear_folder_paths, monkeypatch, tmp_path165):166    """167    Test that when 'base_path' is a relative path in the YAML, it is joined to the YAML file directory, and then168    the items in the config are correctly converted to absolute paths.169    """170    sub_folder = "./my_rel_base"171    config_data = {172        "some_model_folder": {173            "base_path": sub_folder,174            "is_default": True,175            "checkpoints": "checkpoints",176            "some_key": "some_value"177        }178    }179    mock_yaml_load.return_value = config_data180 181    dummy_yaml_name = "dummy_file.yaml"182 183    def fake_abspath(path):184        if path == dummy_yaml_name:185            # If it's the YAML path, treat it like it lives in tmp_path186            return os.path.join(str(tmp_path), dummy_yaml_name)187        return os.path.join(str(tmp_path), path)  # Otherwise, do a normal join relative to tmp_path188 189    def fake_dirname(path):190        # We expect path to be the result of fake_abspath(dummy_yaml_name)191        if path.endswith(dummy_yaml_name):192            return str(tmp_path)193        return os.path.dirname(path)194 195    monkeypatch.setattr(os.path, "abspath", fake_abspath)196    monkeypatch.setattr(os.path, "dirname", fake_dirname)197 198    load_extra_path_config(dummy_yaml_name)199 200    expected_checkpoints = os.path.abspath(os.path.join(str(tmp_path), "my_rel_base", "checkpoints"))201    expected_some_value = os.path.abspath(os.path.join(str(tmp_path), "my_rel_base", "some_value"))202 203    actual_paths = folder_paths.folder_names_and_paths["checkpoints"][0]204    assert len(actual_paths) == 1, "Should have one path added for 'checkpoints'."205    assert actual_paths[0] == expected_checkpoints206 207    actual_paths = folder_paths.folder_names_and_paths["some_key"][0]208    assert len(actual_paths) == 1, "Should have one path added for 'some_key'."209    assert actual_paths[0] == expected_some_value210 211 212@patch("builtins.open", new_callable=mock_open, read_data="dummy yaml content")213@patch("yaml.safe_load")214def test_load_extra_path_config_absolute_base_path(215    mock_yaml_load, _mock_file, clear_folder_paths, monkeypatch, tmp_path216):217    """218    Test that when 'base_path' is an absolute path, each subdirectory is joined with that absolute path,219    rather than being relative to the YAML's directory.220    """221    abs_base = os.path.join(str(tmp_path), "abs_base")222    config_data = {223        "some_absolute_folder": {224            "base_path": abs_base,   # <-- absolute225            "is_default": True,226            "loras": "loras_folder",227            "embeddings": "embeddings_folder"228        }229    }230    mock_yaml_load.return_value = config_data231 232    dummy_yaml_name = "dummy_abs.yaml"233 234    def fake_abspath(path):235        if path == dummy_yaml_name:236            # If it's the YAML path, treat it like it is in tmp_path237            return os.path.join(str(tmp_path), dummy_yaml_name)238        return path  # For absolute base, we just return path directly239 240    def fake_dirname(path):241        return str(tmp_path) if path.endswith(dummy_yaml_name) else os.path.dirname(path)242 243    monkeypatch.setattr(os.path, "abspath", fake_abspath)244    monkeypatch.setattr(os.path, "dirname", fake_dirname)245 246    load_extra_path_config(dummy_yaml_name)247 248    # Expect the final paths to be <abs_base>/loras_folder and <abs_base>/embeddings_folder249    expected_loras = os.path.join(abs_base, "loras_folder")250    expected_embeddings = os.path.join(abs_base, "embeddings_folder")251 252    actual_loras = folder_paths.folder_names_and_paths["loras"][0]253    assert len(actual_loras) == 1, "Should have one path for 'loras'."254    assert actual_loras[0] == os.path.abspath(expected_loras)255 256    actual_embeddings = folder_paths.folder_names_and_paths["embeddings"][0]257    assert len(actual_embeddings) == 1, "Should have one path for 'embeddings'."258    assert actual_embeddings[0] == os.path.abspath(expected_embeddings)259 260 261@patch("builtins.open", new_callable=mock_open, read_data="dummy yaml content")262@patch("yaml.safe_load")263def test_load_extra_path_config_no_base_path(264    mock_yaml_load, _mock_file, clear_folder_paths, monkeypatch, tmp_path265):266    """267    Test that if 'base_path' is not present, each path is joined268    with the directory of the YAML file (unless it's already absolute).269    """270    config_data = {271        "some_folder_without_base": {272            "is_default": True,273            "text_encoders": "clip",274            "diffusion_models": "unet"275        }276    }277    mock_yaml_load.return_value = config_data278 279    dummy_yaml_name = "dummy_no_base.yaml"280 281    def fake_abspath(path):282        if path == dummy_yaml_name:283            return os.path.join(str(tmp_path), dummy_yaml_name)284        return os.path.join(str(tmp_path), path)285 286    def fake_dirname(path):287        return str(tmp_path) if path.endswith(dummy_yaml_name) else os.path.dirname(path)288 289    monkeypatch.setattr(os.path, "abspath", fake_abspath)290    monkeypatch.setattr(os.path, "dirname", fake_dirname)291 292    load_extra_path_config(dummy_yaml_name)293 294    expected_clip = os.path.join(str(tmp_path), "clip")295    expected_unet = os.path.join(str(tmp_path), "unet")296 297    actual_text_encoders = folder_paths.folder_names_and_paths["text_encoders"][0]298    assert len(actual_text_encoders) == 1, "Should have one path for 'text_encoders'."299    assert actual_text_encoders[0] == os.path.abspath(expected_clip)300 301    actual_diffusion = folder_paths.folder_names_and_paths["diffusion_models"][0]302    assert len(actual_diffusion) == 1, "Should have one path for 'diffusion_models'."303    assert actual_diffusion[0] == os.path.abspath(expected_unet)304