diff --git a/tests-unit/utils/extra_config_test.py b/tests-unit/utils/extra_config_test.py index 143e1f2e..b23f5bd0 100644 --- a/tests-unit/utils/extra_config_test.py +++ b/tests-unit/utils/extra_config_test.py @@ -1,11 +1,22 @@ import pytest import yaml import os +import sys from unittest.mock import Mock, patch, mock_open from utils.extra_config import load_extra_path_config import folder_paths + +@pytest.fixture() +def clear_folder_paths(): + # Clear the global dictionary before each test to ensure isolation + original = folder_paths.folder_names_and_paths.copy() + folder_paths.folder_names_and_paths.clear() + yield + folder_paths.folder_names_and_paths = original + + @pytest.fixture def mock_yaml_content(): return { @@ -15,10 +26,12 @@ def mock_yaml_content(): } } + @pytest.fixture def mock_expanded_home(): return '/home/user' + @pytest.fixture def yaml_config_with_appdata(): return """ @@ -27,20 +40,33 @@ def yaml_config_with_appdata(): checkpoints: 'models/checkpoints' """ + @pytest.fixture def mock_yaml_content_appdata(yaml_config_with_appdata): return yaml.safe_load(yaml_config_with_appdata) + @pytest.fixture def mock_expandvars_appdata(): mock = Mock() - mock.side_effect = lambda path: path.replace('%APPDATA%', 'C:/Users/TestUser/AppData/Roaming') + + def expandvars(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.side_effect = expandvars return mock + @pytest.fixture def mock_add_model_folder_path(): return Mock() + @pytest.fixture def mock_expanduser(mock_expanded_home): def _expanduser(path): @@ -49,10 +75,12 @@ def mock_expanduser(mock_expanded_home): return path return _expanduser + @pytest.fixture def mock_yaml_safe_load(mock_yaml_content): return Mock(return_value=mock_yaml_content) + @patch('builtins.open', new_callable=mock_open, read_data="dummy file content") def test_load_extra_model_paths_expands_userpath( mock_file, @@ -88,6 +116,7 @@ def test_load_extra_model_paths_expands_userpath( # Check if open was called with the correct file path mock_file.assert_called_once_with(dummy_yaml_file_name, 'r') + @patch('builtins.open', new_callable=mock_open) def test_load_extra_model_paths_expands_appdata( mock_file, @@ -111,7 +140,10 @@ def test_load_extra_model_paths_expands_appdata( dummy_yaml_file_name = 'dummy_path.yaml' load_extra_path_config(dummy_yaml_file_name) - expected_base_path = 'C:/Users/TestUser/AppData/Roaming/ComfyUI' + if sys.platform == "win32": + expected_base_path = 'C:/Users/TestUser/AppData/Roaming/ComfyUI' + else: + expected_base_path = '/Users/TestUser/AppData/Roaming/ComfyUI' expected_calls = [ ('checkpoints', os.path.join(expected_base_path, 'models/checkpoints'), False), ] @@ -124,3 +156,148 @@ def test_load_extra_model_paths_expands_appdata( # Verify that expandvars was called assert mock_expandvars_appdata.called + + +@patch("builtins.open", new_callable=mock_open, read_data="dummy yaml content") +@patch("yaml.safe_load") +def test_load_extra_path_config_relative_base_path( + mock_yaml_load, _mock_file, clear_folder_paths, monkeypatch, tmp_path +): + """ + Test that when 'base_path' is a relative path in the YAML, it is joined to the YAML file directory, and then + the items in the config are correctly converted to absolute paths. + """ + sub_folder = "./my_rel_base" + config_data = { + "some_model_folder": { + "base_path": sub_folder, + "is_default": True, + "checkpoints": "checkpoints", + "some_key": "some_value" + } + } + mock_yaml_load.return_value = config_data + + dummy_yaml_name = "dummy_file.yaml" + + def fake_abspath(path): + if path == dummy_yaml_name: + # If it's the YAML path, treat it like it lives in tmp_path + return os.path.join(str(tmp_path), dummy_yaml_name) + return os.path.join(str(tmp_path), path) # Otherwise, do a normal join relative to tmp_path + + def fake_dirname(path): + # We expect path to be the result of fake_abspath(dummy_yaml_name) + if path.endswith(dummy_yaml_name): + return str(tmp_path) + return os.path.dirname(path) + + monkeypatch.setattr(os.path, "abspath", fake_abspath) + monkeypatch.setattr(os.path, "dirname", fake_dirname) + + load_extra_path_config(dummy_yaml_name) + + expected_checkpoints = os.path.abspath(os.path.join(str(tmp_path), sub_folder, "checkpoints")) + expected_some_value = os.path.abspath(os.path.join(str(tmp_path), sub_folder, "some_value")) + + actual_paths = folder_paths.folder_names_and_paths["checkpoints"][0] + assert len(actual_paths) == 1, "Should have one path added for 'checkpoints'." + assert actual_paths[0] == expected_checkpoints + + actual_paths = folder_paths.folder_names_and_paths["some_key"][0] + assert len(actual_paths) == 1, "Should have one path added for 'some_key'." + assert actual_paths[0] == expected_some_value + + +@patch("builtins.open", new_callable=mock_open, read_data="dummy yaml content") +@patch("yaml.safe_load") +def test_load_extra_path_config_absolute_base_path( + mock_yaml_load, _mock_file, clear_folder_paths, monkeypatch, tmp_path +): + """ + Test that when 'base_path' is an absolute path, each subdirectory is joined with that absolute path, + rather than being relative to the YAML's directory. + """ + abs_base = os.path.join(str(tmp_path), "abs_base") + config_data = { + "some_absolute_folder": { + "base_path": abs_base, # <-- absolute + "is_default": True, + "loras": "loras_folder", + "embeddings": "embeddings_folder" + } + } + mock_yaml_load.return_value = config_data + + dummy_yaml_name = "dummy_abs.yaml" + + def fake_abspath(path): + if path == dummy_yaml_name: + # If it's the YAML path, treat it like it is in tmp_path + return os.path.join(str(tmp_path), dummy_yaml_name) + return path # For absolute base, we just return path directly + + def fake_dirname(path): + return str(tmp_path) if path.endswith(dummy_yaml_name) else os.path.dirname(path) + + monkeypatch.setattr(os.path, "abspath", fake_abspath) + monkeypatch.setattr(os.path, "dirname", fake_dirname) + + load_extra_path_config(dummy_yaml_name) + + # Expect the final paths to be /loras_folder and /embeddings_folder + expected_loras = os.path.join(abs_base, "loras_folder") + expected_embeddings = os.path.join(abs_base, "embeddings_folder") + + actual_loras = folder_paths.folder_names_and_paths["loras"][0] + assert len(actual_loras) == 1, "Should have one path for 'loras'." + assert actual_loras[0] == os.path.abspath(expected_loras) + + actual_embeddings = folder_paths.folder_names_and_paths["embeddings"][0] + assert len(actual_embeddings) == 1, "Should have one path for 'embeddings'." + assert actual_embeddings[0] == os.path.abspath(expected_embeddings) + + +@patch("builtins.open", new_callable=mock_open, read_data="dummy yaml content") +@patch("yaml.safe_load") +def test_load_extra_path_config_no_base_path( + mock_yaml_load, _mock_file, clear_folder_paths, monkeypatch, tmp_path +): + """ + Test that if 'base_path' is not present, each path is joined + with the directory of the YAML file (unless it's already absolute). + """ + config_data = { + "some_folder_without_base": { + "is_default": True, + "text_encoders": "clip", + "diffusion_models": "unet" + } + } + mock_yaml_load.return_value = config_data + + dummy_yaml_name = "dummy_no_base.yaml" + + def fake_abspath(path): + if path == dummy_yaml_name: + return os.path.join(str(tmp_path), dummy_yaml_name) + return os.path.join(str(tmp_path), path) + + def fake_dirname(path): + return str(tmp_path) if path.endswith(dummy_yaml_name) else os.path.dirname(path) + + monkeypatch.setattr(os.path, "abspath", fake_abspath) + monkeypatch.setattr(os.path, "dirname", fake_dirname) + + load_extra_path_config(dummy_yaml_name) + + expected_clip = os.path.join(str(tmp_path), "clip") + expected_unet = os.path.join(str(tmp_path), "unet") + + actual_text_encoders = folder_paths.folder_names_and_paths["text_encoders"][0] + assert len(actual_text_encoders) == 1, "Should have one path for 'text_encoders'." + assert actual_text_encoders[0] == os.path.abspath(expected_clip) + + actual_diffusion = folder_paths.folder_names_and_paths["diffusion_models"][0] + assert len(actual_diffusion) == 1, "Should have one path for 'diffusion_models'." + assert actual_diffusion[0] == os.path.abspath(expected_unet) diff --git a/utils/extra_config.py b/utils/extra_config.py index 415db042..d7b59285 100644 --- a/utils/extra_config.py +++ b/utils/extra_config.py @@ -6,6 +6,7 @@ import logging def load_extra_path_config(yaml_path): with open(yaml_path, 'r') as stream: config = yaml.safe_load(stream) + yaml_dir = os.path.dirname(os.path.abspath(yaml_path)) for c in config: conf = config[c] if conf is None: @@ -14,6 +15,8 @@ def load_extra_path_config(yaml_path): if "base_path" in conf: base_path = conf.pop("base_path") base_path = os.path.expandvars(os.path.expanduser(base_path)) + if not os.path.isabs(base_path): + base_path = os.path.abspath(os.path.join(yaml_dir, base_path)) is_default = False if "is_default" in conf: is_default = conf.pop("is_default") @@ -22,10 +25,9 @@ def load_extra_path_config(yaml_path): if len(y) == 0: continue full_path = y - if base_path is not None: + if base_path: full_path = os.path.join(base_path, full_path) elif not os.path.isabs(full_path): - yaml_dir = os.path.dirname(os.path.abspath(yaml_path)) full_path = os.path.abspath(os.path.join(yaml_dir, y)) logging.info("Adding extra search path {} {}".format(x, full_path)) folder_paths.add_model_folder_path(x, full_path, is_default)