Add unit test for expanding base path.

tests-unit/main_test.py

@@ -1,46 +1,67 @@
 import pytest
-from unittest.mock import patch, mock_open
 import os
+import logging
+from unittest.mock import patch
 
+# Import the function we're testing
 from main import load_extra_path_config
 
 @pytest.fixture
-def mock_add_path():
+def mock_yaml_file():
-    with patch('folder_paths.add_model_folder_path') as mock:
+    yaml_content = """
-        yield mock
+    test_config:
-
+        base_path: ~/App/
-@pytest.mark.parametrize("os_type,base_path,expected_paths", [
+        checkpoint:
-    ("windows", "~\\AppData\\Local\\ComfyUI", {
+            subfolder1
-        'checkpoints': os.path.expanduser('~\\AppData\\Local\\ComfyUI\\models\\checkpoints'),
+            subfolder2
-        'clip': os.path.expanduser('~\\AppData\\Local\\ComfyUI\\models\\clip')
+        lora: otherfolder
-    }),
-    ("linux", "~/comfyui", {
-        'checkpoints': os.path.expanduser('~/comfyui/models/checkpoints'),
-        'clip': os.path.expanduser('~/comfyui/models/clip')
-    }),
-    ("mac", "~/Library/Application Support/ComfyUI", {
-        'checkpoints': os.path.expanduser('~/Library/Application Support/ComfyUI/models/checkpoints'),
-        'clip': os.path.expanduser('~/Library/Application Support/ComfyUI/models/clip')
-    }),
-])
-def test_load_config_with_base_path(mock_add_path, os_type, base_path, expected_paths, monkeypatch):
-    if os_type == "windows":
-        monkeypatch.setattr(os, 'name', 'nt')
-        monkeypatch.setattr(os.path, 'sep', '\\')
-    else:
-        monkeypatch.setattr(os, 'name', 'posix')
-        monkeypatch.setattr(os.path, 'sep', '/')
-
-    yaml_content = f"""
-    comfyui:
-        base_path: "{base_path}"
-        checkpoints: models{os.path.sep}checkpoints{os.path.sep}
-        clip: models{os.path.sep}clip{os.path.sep}
     """
-    mock_open_file = mock_open(read_data=yaml_content)
+    return yaml_content
 
-    with patch('builtins.open', mock_open_file):
+@pytest.fixture
-        load_extra_path_config('dummy_path.yaml')
+def mock_expanded_home():
+    return '/home/user'
 
-    for model_type, expected_path in expected_paths.items():
+@patch('os.path.expanduser')
-        mock_add_path.assert_any_call(model_type, expected_path)
+@patch('folder_paths.add_model_folder_path')
+def test_load_extra_path_config(mock_add_model_folder_path, mock_expanduser, mock_yaml_file, mock_expanded_home, tmp_path):
+    # Setup
+    mock_expanduser.return_value = os.path.join(mock_expanded_home, 'App')
+    yaml_path = tmp_path / "test_config.yaml"
+    with open(yaml_path, 'w') as f:
+        f.write(mock_yaml_file)
+
+    # Call the function
+    load_extra_path_config(yaml_path)
+
+    # Assertions
+    expected_calls = [
+        ('checkpoint', os.path.join(mock_expanded_home, 'App', 'subfolder1 subfolder2')),
+        ('lora', os.path.join(mock_expanded_home, 'App', 'otherfolder'))
+    ]
+
+    assert mock_add_model_folder_path.call_count == len(expected_calls)
+    for call in mock_add_model_folder_path.call_args_list:
+        assert call.args in expected_calls
+
+    # Check if expanduser was called with the correct path
+    mock_expanduser.assert_called_once_with('~/App/')
+
+@pytest.fixture
+def caplog(caplog):
+    caplog.set_level(logging.INFO)
+    return caplog
+
+def test_load_extra_path_config_logging(mock_yaml_file, tmp_path, caplog):
+    # Setup
+    yaml_path = tmp_path / "test_config.yaml"
+    with open(yaml_path, 'w') as f:
+        f.write(mock_yaml_file)
+
+    # Call the function
+    with patch('folder_paths.add_model_folder_path'):
+        load_extra_path_config(yaml_path)
+
+    # Check logged messages
+    assert "Adding extra search path checkpoint " in caplog.text
+    assert "Adding extra search path lora " in caplog.text