From cd2629dd6767f97d40b7e20d1d1a6a98f090a991 Mon Sep 17 00:00:00 2001
From: Robin Huang <robin.j.huang@gmail.com>
Date: Mon, 9 Sep 2024 16:08:54 +0900
Subject: [PATCH] Add unit test for expanding base path.

---
 tests-unit/main_test.py | 95 +++++++++++++++++++++++++----------------
 1 file changed, 58 insertions(+), 37 deletions(-)

diff --git a/tests-unit/main_test.py b/tests-unit/main_test.py
index da730e6ba..8384266a9 100644
--- a/tests-unit/main_test.py
+++ b/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():
-    with patch('folder_paths.add_model_folder_path') as mock:
-        yield mock
-
-@pytest.mark.parametrize("os_type,base_path,expected_paths", [
-    ("windows", "~\\AppData\\Local\\ComfyUI", {
-        'checkpoints': os.path.expanduser('~\\AppData\\Local\\ComfyUI\\models\\checkpoints'),
-        'clip': os.path.expanduser('~\\AppData\\Local\\ComfyUI\\models\\clip')
-    }),
-    ("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}
+def mock_yaml_file():
+    yaml_content = """
+    test_config:
+        base_path: ~/App/
+        checkpoint:
+            subfolder1
+            subfolder2
+        lora: otherfolder
     """
-    mock_open_file = mock_open(read_data=yaml_content)
+    return yaml_content
 
-    with patch('builtins.open', mock_open_file):
-        load_extra_path_config('dummy_path.yaml')
+@pytest.fixture
+def mock_expanded_home():
+    return '/home/user'
 
-    for model_type, expected_path in expected_paths.items():
-        mock_add_path.assert_any_call(model_type, expected_path)
+@patch('os.path.expanduser')
+@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