Expand variables in base_path for extra_config_paths.yaml. (#4893)

* Expand variables in base_path for extra_config_paths.yaml.

* Fix comments.
This commit is contained in:
Robin Huang 2024-09-12 14:52:06 +09:00 committed by GitHub
parent 9f4daca9d9
commit d247bc5a9c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 56 additions and 1 deletions

View File

@ -19,6 +19,24 @@ def mock_yaml_content():
def mock_expanded_home(): def mock_expanded_home():
return '/home/user' return '/home/user'
@pytest.fixture
def yaml_config_with_appdata():
return """
test_config:
base_path: '%APPDATA%/ComfyUI'
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')
return mock
@pytest.fixture @pytest.fixture
def mock_add_model_folder_path(): def mock_add_model_folder_path():
return Mock() return Mock()
@ -67,3 +85,40 @@ def test_load_extra_model_paths_expands_userpath(
# Check if open was called with the correct file path # Check if open was called with the correct file path
mock_file.assert_called_once_with(dummy_yaml_file_name, 'r') 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,
monkeypatch,
mock_add_model_folder_path,
mock_expandvars_appdata,
yaml_config_with_appdata,
mock_yaml_content_appdata
):
# Set the mock_file to return yaml with appdata as a variable
mock_file.return_value.read.return_value = yaml_config_with_appdata
# Attach mocks
monkeypatch.setattr(folder_paths, 'add_model_folder_path', mock_add_model_folder_path)
monkeypatch.setattr(os.path, 'expandvars', mock_expandvars_appdata)
monkeypatch.setattr(yaml, 'safe_load', Mock(return_value=mock_yaml_content_appdata))
# Mock expanduser to do nothing (since we're not testing it here)
monkeypatch.setattr(os.path, 'expanduser', lambda x: x)
dummy_yaml_file_name = 'dummy_path.yaml'
load_extra_path_config(dummy_yaml_file_name)
expected_base_path = 'C:/Users/TestUser/AppData/Roaming/ComfyUI'
expected_calls = [
('checkpoints', os.path.join(expected_base_path, 'models/checkpoints')),
]
assert mock_add_model_folder_path.call_count == len(expected_calls)
# Check the base path variable was expanded
for actual_call, expected_call in zip(mock_add_model_folder_path.call_args_list, expected_calls):
assert actual_call.args == expected_call
# Verify that expandvars was called
assert mock_expandvars_appdata.called

View File

@ -13,7 +13,7 @@ def load_extra_path_config(yaml_path):
base_path = None base_path = None
if "base_path" in conf: if "base_path" in conf:
base_path = conf.pop("base_path") base_path = conf.pop("base_path")
base_path = os.path.expanduser(base_path) base_path = os.path.expandvars(os.path.expanduser(base_path))
for x in conf: for x in conf:
for y in conf[x].split("\n"): for y in conf[x].split("\n"):
if len(y) == 0: if len(y) == 0: