mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2025-01-11 02:15:17 +00:00
fixed a bug where a relative path was not converted to a full path
Signed-off-by: bigcat88 <bigcat88@icloud.com>
This commit is contained in:
parent
d0f3752e33
commit
126cdc582b
@ -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 <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,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)
|
||||
|
Loading…
Reference in New Issue
Block a user