Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

(fix): load_extra_path_config: relative path not converted to a full path #6395

Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
181 changes: 179 additions & 2 deletions tests-unit/utils/extra_config_test.py
Original file line number Diff line number Diff line change
@@ -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 {
Expand All @@ -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 """
Expand All @@ -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):
Expand All @@ -49,10 +75,12 @@ def _expanduser(path):
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,
Expand Down Expand Up @@ -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,
Expand All @@ -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),
]
Expand All @@ -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 <abs_base>/loras_folder and <abs_base>/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)
6 changes: 4 additions & 2 deletions utils/extra_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
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:
Expand All @@ -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")
Expand All @@ -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)
Loading