Vito Sansevero
4400444619
test_image_monitor_singleton replaced utils.comfyui_integration in sys.modules with a MagicMock and never restored it, so running it before test_comfyui_integration failed eight tests. The mock was unused; drop it and install the server stub with setdefault in every owned test module. The PromptSearchList IS_CHANGED test slept 10 ms between two time.time() reads; the Windows clock ticks every 15.6 ms, so it was flaky there. Mock the clock with an incrementing counter instead.
270 lines
10 KiB
Python
270 lines
10 KiB
Python
"""Tests for gallery root validation (GalleryConfig.validate_gallery_root).
|
|
|
|
A gallery root may only point inside one of ComfyUI's own directories
|
|
(output/input/temp/user, as reported by ``folder_paths``) or inside a
|
|
directory listed in ``PROMPT_MANAGER_EXTRA_GALLERY_ROOTS``. Anything else,
|
|
including filesystem roots and the home directory, is rejected so the
|
|
image-serving routes can never be pointed at arbitrary files.
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
import types
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import MagicMock
|
|
|
|
_mock_server = MagicMock()
|
|
_mock_server.PromptServer.instance.routes = MagicMock()
|
|
sys.modules.setdefault("server", _mock_server)
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
from py.config import GalleryConfig # noqa: E402
|
|
|
|
EXTRA_ROOTS_ENV = "PROMPT_MANAGER_EXTRA_GALLERY_ROOTS"
|
|
|
|
|
|
class GalleryRootTestCase(unittest.TestCase):
|
|
"""Shared fixture: a fake ComfyUI tree exposed through a stub folder_paths."""
|
|
|
|
def setUp(self):
|
|
self.tmpdir = tempfile.mkdtemp()
|
|
self.addCleanup(self._rmtree, self.tmpdir)
|
|
|
|
self.comfy_dir = Path(self.tmpdir) / "ComfyUI"
|
|
self.output_dir = self.comfy_dir / "output"
|
|
self.input_dir = self.comfy_dir / "input"
|
|
self.output_dir.mkdir(parents=True)
|
|
self.input_dir.mkdir(parents=True)
|
|
|
|
self._orig_folder_paths = sys.modules.get("folder_paths")
|
|
sys.modules["folder_paths"] = types.SimpleNamespace(
|
|
base_path=str(self.comfy_dir),
|
|
get_output_directory=lambda: str(self.output_dir),
|
|
get_input_directory=lambda: str(self.input_dir),
|
|
get_temp_directory=self._raise_runtime_error,
|
|
# get_user_directory intentionally absent
|
|
)
|
|
self.addCleanup(self._restore_folder_paths)
|
|
|
|
self._orig_extra = os.environ.pop(EXTRA_ROOTS_ENV, None)
|
|
self.addCleanup(self._restore_extra_env)
|
|
|
|
@staticmethod
|
|
def _raise_runtime_error():
|
|
raise RuntimeError("temp directory unavailable")
|
|
|
|
@staticmethod
|
|
def _rmtree(path):
|
|
import shutil
|
|
|
|
shutil.rmtree(path, ignore_errors=True)
|
|
|
|
def _restore_folder_paths(self):
|
|
if self._orig_folder_paths is None:
|
|
sys.modules.pop("folder_paths", None)
|
|
else:
|
|
sys.modules["folder_paths"] = self._orig_folder_paths
|
|
|
|
def _restore_extra_env(self):
|
|
if self._orig_extra is None:
|
|
os.environ.pop(EXTRA_ROOTS_ENV, None)
|
|
else:
|
|
os.environ[EXTRA_ROOTS_ENV] = self._orig_extra
|
|
|
|
|
|
class TestValidateGalleryRootRejects(GalleryRootTestCase):
|
|
|
|
def test_filesystem_root_rejected(self):
|
|
ok, reason = GalleryConfig.validate_gallery_root(os.path.abspath(os.sep))
|
|
self.assertFalse(ok)
|
|
self.assertIn("root", reason.lower())
|
|
|
|
@unittest.skipUnless(os.name == "nt", "drive roots only exist on Windows")
|
|
def test_windows_drive_root_rejected(self):
|
|
ok, _ = GalleryConfig.validate_gallery_root("C:\\")
|
|
self.assertFalse(ok)
|
|
|
|
def test_home_directory_rejected(self):
|
|
ok, reason = GalleryConfig.validate_gallery_root(os.path.expanduser("~"))
|
|
self.assertFalse(ok)
|
|
self.assertIn("home", reason.lower())
|
|
|
|
def test_directory_outside_comfyui_rejected(self):
|
|
outside = Path(self.tmpdir) / "elsewhere"
|
|
outside.mkdir()
|
|
ok, _ = GalleryConfig.validate_gallery_root(str(outside))
|
|
self.assertFalse(ok)
|
|
|
|
def test_nonexistent_path_rejected(self):
|
|
missing = self.output_dir / "does_not_exist"
|
|
ok, reason = GalleryConfig.validate_gallery_root(str(missing))
|
|
self.assertFalse(ok)
|
|
self.assertIn("exist", reason.lower())
|
|
|
|
def test_file_rejected(self):
|
|
file_path = self.output_dir / "image.png"
|
|
file_path.write_bytes(b"PNG")
|
|
ok, reason = GalleryConfig.validate_gallery_root(str(file_path))
|
|
self.assertFalse(ok)
|
|
self.assertIn("directory", reason.lower())
|
|
|
|
def test_empty_and_non_string_rejected(self):
|
|
self.assertFalse(GalleryConfig.validate_gallery_root("")[0])
|
|
self.assertFalse(GalleryConfig.validate_gallery_root(" ")[0])
|
|
self.assertFalse(GalleryConfig.validate_gallery_root(None)[0])
|
|
self.assertFalse(GalleryConfig.validate_gallery_root(42)[0])
|
|
|
|
def test_sibling_with_shared_prefix_rejected(self):
|
|
"""'output2' must not pass because it string-starts with 'output'."""
|
|
sibling = self.comfy_dir / "output2"
|
|
sibling.mkdir()
|
|
ok, _ = GalleryConfig.validate_gallery_root(str(sibling))
|
|
self.assertFalse(ok)
|
|
|
|
def test_traversal_out_of_output_rejected(self):
|
|
outside = Path(self.tmpdir) / "elsewhere"
|
|
outside.mkdir()
|
|
sneaky = os.path.join(str(self.output_dir), "..", "..", "elsewhere")
|
|
ok, _ = GalleryConfig.validate_gallery_root(sneaky)
|
|
self.assertFalse(ok)
|
|
|
|
def test_symlink_escaping_output_rejected(self):
|
|
outside = Path(self.tmpdir) / "elsewhere"
|
|
outside.mkdir()
|
|
link = self.output_dir / "escape"
|
|
try:
|
|
os.symlink(str(outside), str(link), target_is_directory=True)
|
|
except (OSError, NotImplementedError):
|
|
self.skipTest("symlinks not available on this platform")
|
|
ok, _ = GalleryConfig.validate_gallery_root(str(link))
|
|
self.assertFalse(ok)
|
|
|
|
def test_env_root_that_is_filesystem_root_is_ignored(self):
|
|
"""Listing '/' in the env var must not re-open the whole disk."""
|
|
os.environ[EXTRA_ROOTS_ENV] = os.path.abspath(os.sep)
|
|
outside = Path(self.tmpdir) / "elsewhere"
|
|
outside.mkdir()
|
|
ok, _ = GalleryConfig.validate_gallery_root(str(outside))
|
|
self.assertFalse(ok)
|
|
|
|
|
|
class TestValidateGalleryRootAccepts(GalleryRootTestCase):
|
|
|
|
def test_output_directory_itself_accepted(self):
|
|
ok, reason = GalleryConfig.validate_gallery_root(str(self.output_dir))
|
|
self.assertTrue(ok, reason)
|
|
|
|
def test_subdirectory_of_output_accepted(self):
|
|
sub = self.output_dir / "renders" / "2026"
|
|
sub.mkdir(parents=True)
|
|
ok, reason = GalleryConfig.validate_gallery_root(str(sub))
|
|
self.assertTrue(ok, reason)
|
|
|
|
def test_subdirectory_of_input_accepted(self):
|
|
sub = self.input_dir / "refs"
|
|
sub.mkdir()
|
|
ok, reason = GalleryConfig.validate_gallery_root(str(sub))
|
|
self.assertTrue(ok, reason)
|
|
|
|
def test_directory_listed_in_env_accepted(self):
|
|
extra_a = Path(self.tmpdir) / "extra_a"
|
|
extra_b = Path(self.tmpdir) / "extra_b"
|
|
extra_a.mkdir()
|
|
extra_b.mkdir()
|
|
os.environ[EXTRA_ROOTS_ENV] = os.pathsep.join([str(extra_a), str(extra_b)])
|
|
|
|
ok_a, reason_a = GalleryConfig.validate_gallery_root(str(extra_a))
|
|
ok_b, reason_b = GalleryConfig.validate_gallery_root(str(extra_b / "."))
|
|
self.assertTrue(ok_a, reason_a)
|
|
self.assertTrue(ok_b, reason_b)
|
|
|
|
def test_relative_path_resolves_against_comfyui_base_not_cwd(self):
|
|
sub = self.output_dir / "rel"
|
|
sub.mkdir()
|
|
orig_cwd = os.getcwd()
|
|
os.chdir(self.tmpdir)
|
|
self.addCleanup(os.chdir, orig_cwd)
|
|
|
|
ok, reason = GalleryConfig.validate_gallery_root("output/rel")
|
|
self.assertTrue(ok, reason)
|
|
self.assertEqual(
|
|
GalleryConfig.resolve_gallery_root("output/rel"),
|
|
os.path.normcase(os.path.realpath(sub)),
|
|
)
|
|
# A path that only exists relative to the CWD is not accepted.
|
|
(Path(self.tmpdir) / "rel").mkdir()
|
|
self.assertFalse(GalleryConfig.validate_gallery_root("rel")[0])
|
|
|
|
def test_accepts_without_folder_paths_when_env_lists_dir(self):
|
|
sys.modules.pop("folder_paths", None)
|
|
extra = Path(self.tmpdir) / "extra"
|
|
extra.mkdir()
|
|
os.environ[EXTRA_ROOTS_ENV] = str(extra)
|
|
ok, reason = GalleryConfig.validate_gallery_root(str(extra))
|
|
self.assertTrue(ok, reason)
|
|
|
|
def test_rejects_everything_without_folder_paths_and_env(self):
|
|
sys.modules.pop("folder_paths", None)
|
|
ok, _ = GalleryConfig.validate_gallery_root(str(self.output_dir))
|
|
self.assertFalse(ok)
|
|
|
|
|
|
class TestGetAllOutputDirsFiltersRoots(GalleryRootTestCase):
|
|
"""PromptManagerAPI._get_all_output_dirs() only returns validated roots."""
|
|
|
|
def setUp(self):
|
|
super().setUp()
|
|
from py.api import PromptManagerAPI
|
|
|
|
self.api = PromptManagerAPI()
|
|
self.api._cached_output_dir = None
|
|
self._orig_dirs = list(GalleryConfig.MONITORING_DIRECTORIES)
|
|
self.addCleanup(self._restore_dirs)
|
|
|
|
def _restore_dirs(self):
|
|
GalleryConfig.MONITORING_DIRECTORIES = self._orig_dirs
|
|
|
|
def _dirs(self):
|
|
return [os.path.normcase(str(p)) for p in self.api._get_all_output_dirs()]
|
|
|
|
def test_filesystem_root_is_never_served(self):
|
|
GalleryConfig.MONITORING_DIRECTORIES = [os.path.abspath(os.sep)]
|
|
dirs = self._dirs()
|
|
self.assertNotIn(os.path.normcase(os.path.abspath(os.sep)), dirs)
|
|
self.assertEqual(dirs, [os.path.normcase(os.path.realpath(self.output_dir))])
|
|
|
|
def test_outside_directory_is_dropped_and_valid_one_kept(self):
|
|
outside = Path(self.tmpdir) / "elsewhere"
|
|
outside.mkdir()
|
|
good = self.output_dir / "good"
|
|
good.mkdir()
|
|
GalleryConfig.MONITORING_DIRECTORIES = [str(outside), str(good)]
|
|
self.assertEqual(self._dirs(), [os.path.normcase(os.path.realpath(good))])
|
|
|
|
def test_all_rejected_falls_back_to_comfyui_output_only(self):
|
|
outside = Path(self.tmpdir) / "elsewhere"
|
|
outside.mkdir()
|
|
GalleryConfig.MONITORING_DIRECTORIES = [str(outside)]
|
|
self.assertEqual(
|
|
self._dirs(), [os.path.normcase(os.path.realpath(self.output_dir))]
|
|
)
|
|
|
|
def test_no_configured_dirs_uses_auto_detection(self):
|
|
GalleryConfig.MONITORING_DIRECTORIES = []
|
|
self.api._find_comfyui_output_dir = lambda: str(self.output_dir)
|
|
self.assertEqual(self._dirs(), [os.path.normcase(str(self.output_dir))])
|
|
|
|
def test_nothing_available_returns_empty_list(self):
|
|
sys.modules.pop("folder_paths", None)
|
|
outside = Path(self.tmpdir) / "elsewhere"
|
|
outside.mkdir()
|
|
GalleryConfig.MONITORING_DIRECTORIES = [str(outside)]
|
|
self.assertEqual(self._dirs(), [])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|