Automated test for #30

comfy_compatibility/workspace.py

@@ -37,48 +37,65 @@ def auto_patch_workspace_and_restart():
     ]
 
     patched_any_file = False
+    files_to_ignore = []
 
     for dir_name in target_base_dirs:
         start_dir = os.path.join(workspace_dir, dir_name)
         if not os.path.isdir(start_dir):
             continue
 
+        dirs_with_py_files = set()
         for dirpath, _, filenames in os.walk(start_dir):
-            init_py_path = os.path.join(dirpath, '__init__.py')
-
-            if os.path.exists(init_py_path):
-                continue
-
             if any(fname.endswith('.py') for fname in filenames):
+                dirs_with_py_files.add(dirpath)
+
+        if not dirs_with_py_files:
+            continue
+
+        dirs_to_initialize = set()
+        for dirpath in dirs_with_py_files:
+            parent = dirpath
+            while len(parent) >= len(start_dir):
+                dirs_to_initialize.add(parent)
+                new_parent = os.path.dirname(parent)
+                if new_parent == parent:
+                    break
+                parent = new_parent
+
+        for dirpath in sorted(list(dirs_to_initialize)):
+            init_py_path = os.path.join(dirpath, '__init__.py')
+            if not os.path.exists(init_py_path):
                 logger.debug(f"  Initializing package: {dirpath}")
                 try:
                     with open(init_py_path, 'w') as f:
                         pass
                     patched_any_file = True
-
+                    files_to_ignore.append(init_py_path)
-                    if git_is_available:
-                        try:
-                            relative_path = os.path.relpath(init_py_path, workspace_dir).replace(os.sep, '/')
-
-                            exclude_file = os.path.join(workspace_dir, '.git', 'info', 'exclude')
-                            os.makedirs(os.path.dirname(exclude_file), exist_ok=True)
-
-                            content = ""
-                            if os.path.exists(exclude_file):
-                                with open(exclude_file, 'r') as f_read:
-                                    content = f_read.read()
-
-                            if relative_path not in content.splitlines():
-                                with open(exclude_file, 'a') as f_append:
-                                    f_append.write(f"\n{relative_path}")
-                                logger.debug(f"  Ignoring via .git/info/exclude: {relative_path}")
-
-                        except Exception as e:
-                            logger.debug(f"Warning: Could not add {relative_path} to Git exclude file. Error: {e}")
-
                 except OSError as e:
                     logger.debug(f"Warning: Could not create {init_py_path}. Error: {e}")
 
+    if git_is_available and files_to_ignore:
+        try:
+            exclude_file = os.path.join(workspace_dir, '.git', 'info', 'exclude')
+            os.makedirs(os.path.dirname(exclude_file), exist_ok=True)
+
+            existing_lines = set()
+            if os.path.exists(exclude_file):
+                with open(exclude_file, 'r') as f_read:
+                    existing_lines = set(line.strip() for line in f_read)
+
+            with open(exclude_file, 'a') as f_append:
+                for init_py_path in files_to_ignore:
+                    relative_path = os.path.relpath(init_py_path, workspace_dir).replace(os.sep, '/')
+                    if relative_path not in existing_lines:
+                        f_append.write(f"\n{relative_path}")
+                        existing_lines.add(relative_path)
+                        logger.debug(f"  Ignoring via .git/info/exclude: {relative_path}")
+
+        except Exception as e:
+            logger.debug(f"Warning: Could not update Git exclude file. Error: {e}")
+
     if patched_any_file:
         logger.debug("Found and initialized Python package directories in your workspace. This is a one-time operation to enable proper imports. Now restarting...")
         os.execv(sys.executable, [sys.executable] + sys.argv)
+

tests/issues/__test_30_comfy_workspace_conflicts.py

@@ -1,18 +1,191 @@
-"""
+import pytest
-# get the workspace and set it up
+import subprocess
-git clone git@github.com:comfyanonymous/ComfyUI.git comfyanonymous_ComfyUI
+import sys
-cd comfyanonymous_ComfyUI
+import shutil
-uv venv
+import time
-source .venv/bin/activate
+from pathlib import Path
-uv pip install --torch-backend=auto requirements.txt
 
-# install livepeer with up to date comfyui lts
+# Timeout in seconds for the server to start and print the expected lines.
-uv pip install --torch-backend=auto git+https://github.com/livepeer/comfystream.git --overrides=<(echo "comfyui@git+https://github.com/hiddenswitch/ComfyUI.git@fixes/issue-30")
+SERVER_START_TIMEOUT = 20
+# Repository URLs and references
+COMFYUI_REPO = "https://github.com/comfyanonymous/ComfyUI.git"
+COMFYSTREAM_REPO = "https://github.com/doctorpangloss/comfystream.git"
+COMFYSTREAM_COMMIT = "f2f7929def53a4853cc5a1c2774aea70775ce2ff"
+COMFYUI_LTS_REPO = "https://github.com/hiddenswitch/ComfyUI.git"
+COMFYUI_LTS_COMMIT = "6cffd4c4c22e45a8bf07d3a3565ca88de1ca6168"
 
-# install the nodes
-curl -L -o comfystream.zip https://github.com/livepeer/comfystream/archive/refs/heads/main.zip
-mkdir -p custom_nodes
-unzip comfystream.zip -d custom_nodes/
-rm comfystream.zip
 
-"""
+def run_command(cmd, cwd, desc, shell=False):
+    """Helper function to run a command and raise an error if it fails."""
+    print(f"\n--- {desc} ---")
+    log_cmd = cmd if isinstance(cmd, str) else ' '.join(map(str, cmd))
+    print(f"Running command: {log_cmd}")
+    result = subprocess.run(
+        cmd,
+        cwd=cwd,
+        capture_output=True,
+        text=True,
+        check=False,  # We check manually to provide better error logs
+        shell=shell
+    )
+    if result.returncode != 0:
+        print("--- STDOUT ---")
+        print(result.stdout)
+        print("--- STDERR ---")
+        print(result.stderr)
+        pytest.fail(f"Command failed: {desc}", pytrace=False)
+    print(f"--- Success: {desc} ---")
+    return result
+
+
+@pytest.fixture(scope="module")
+def comfyui_workspace(tmp_path_factory):
+    """
+    A pytest fixture that sets up the entire ComfyUI workspace in a temporary
+    directory for a single test module run.
+    """
+    workspace_root = tmp_path_factory.mktemp("comfyui_ws")
+    comfyui_dir = workspace_root / "ComfyUI"
+
+    # 1. Clone the ComfyUI repository
+    run_command(["git", "clone", COMFYUI_REPO, str(comfyui_dir)], cwd=workspace_root, desc="Cloning ComfyUI")
+
+    # 2. Set up the virtual environment using system uv
+    run_command(["uv", "venv"], cwd=comfyui_dir, desc="Creating virtual environment with uv")
+    venv_python = comfyui_dir / ".venv" / ("Scripts" if sys.platform == "win32" else "bin") / "python"
+
+    # Determine activation command and construct shell commands
+    if sys.platform == "win32":
+        # On Windows, we 'call' the activate.bat script. Quotes handle spaces in path.
+        activate_cmd = f'call "{comfyui_dir / ".venv" / "Scripts" / "activate.bat"}"'
+    else:
+        # On Unix-like systems, we '.' (source) the activate script for POSIX compliance.
+        activate_cmd = f'. "{comfyui_dir / ".venv" / "bin" / "activate"}"'
+
+    uv_pip_install_base = "uv pip install --torch-backend=auto"
+
+    # 3. Install base requirements
+    install_reqs_cmd = f"{activate_cmd} && {uv_pip_install_base} -r requirements.txt"
+    run_command(install_reqs_cmd, cwd=comfyui_dir, desc="Installing requirements.txt", shell=True)
+
+    # 4. Install comfystream with the specific comfyui override
+    overrides_content = f"comfyui@git+{COMFYUI_LTS_REPO}@{COMFYUI_LTS_COMMIT}"
+    overrides_file = workspace_root / "overrides.txt"
+    overrides_file.write_text(overrides_content)
+
+    # Using an absolute path for the overrides file is safest for shell execution.
+    install_comfystream_cmd_str = (
+        f'{activate_cmd} && {uv_pip_install_base} '
+        f'git+{COMFYSTREAM_REPO}@{COMFYSTREAM_COMMIT} '
+        f'--overrides={overrides_file.resolve()}'
+    )
+    run_command(install_comfystream_cmd_str, cwd=comfyui_dir, desc="Installing comfystream with overrides", shell=True)
+
+    # 5. Additionally clone comfystream into the custom_nodes directory and check out the specific commit
+    custom_nodes_dir = comfyui_dir / "custom_nodes"
+    comfystream_custom_node_dir = custom_nodes_dir / "comfystream"
+    custom_nodes_dir.mkdir(exist_ok=True)
+    run_command(
+        ["git", "clone", COMFYSTREAM_REPO, str(comfystream_custom_node_dir)],
+        cwd=comfyui_dir,
+        desc="Cloning comfystream into custom_nodes"
+    )
+    run_command(
+        ["git", "checkout", COMFYSTREAM_COMMIT],
+        cwd=comfystream_custom_node_dir,
+        desc=f"Checking out comfystream commit {COMFYSTREAM_COMMIT[:7]}"
+    )
+
+
+    # Yield the necessary paths to the test function
+    yield venv_python, comfyui_dir
+
+    # Teardown is handled automatically by pytest's tmp_path_factory
+
+
+def test_server_starts_with_comfystream(comfyui_workspace):
+    """
+    Tests if the ComfyUI server starts correctly with the comfystream package
+    and prints the expected initialization messages.
+    """
+    venv_python, comfyui_dir = comfyui_workspace
+
+    # --- Assert that the installed package from overrides is structured correctly ---
+    site_packages_cmd = [
+        str(venv_python),
+        "-c",
+        "import sysconfig; print(sysconfig.get_paths()['purelib'])"
+    ]
+    site_packages_result = run_command(site_packages_cmd, cwd=comfyui_dir, desc="Finding site-packages directory")
+    site_packages_path = Path(site_packages_result.stdout.strip())
+    comfy_api_init_path = site_packages_path / "comfy" / "api" / "__init__.py"
+
+    assert comfy_api_init_path.is_file(), (
+        f"The installed comfyui package is missing the api module. "
+        f"Expected to find: {comfy_api_init_path}"
+    )
+    print(f"--- Success: Found {comfy_api_init_path} ---")
+
+    main_py_path = comfyui_dir / "main.py"
+
+    # Start the server as a background process
+    process = None
+    try:
+        print("\n--- Starting ComfyUI Server ---")
+        process = subprocess.Popen(
+            [str(venv_python), str(main_py_path)],
+            cwd=comfyui_dir,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT, # Redirect stderr to stdout
+            text=True,
+            bufsize=1 # Line-buffered
+        )
+
+        output_lines = []
+        found_gui_msg = False
+        found_server_msg = False
+
+        start_time = time.time()
+
+        # Read output line-by-line with a timeout
+        while time.time() - start_time < SERVER_START_TIMEOUT:
+            line = process.stdout.readline()
+            if not line and process.poll() is not None:
+                pytest.fail("Server process terminated unexpectedly.")
+
+            if line:
+                print(line, end='')
+                output_lines.append(line)
+                if "Initializing LocalComfyStreamServer" in line:
+                    found_server_msg = True
+                if "To see the GUI go to: http://127.0.0.1:8188" in line:
+                    found_gui_msg = True
+                    if not found_server_msg:
+                        pytest.fail(
+                            "GUI message appeared before LocalComfyStreamServer was initialized. "
+                            "This indicates a custom node loading problem."
+                        )
+
+            if found_gui_msg and found_server_msg:
+                break
+
+        # Final assertions
+        assert found_gui_msg, "GUI startup message was not found in the output."
+        assert found_server_msg, "LocalComfyStreamServer initialization message was not found."
+
+    finally:
+        if process:
+            print("\n--- Terminating ComfyUI Server ---")
+            process.terminate()
+            try:
+                process.wait(timeout=5)
+            except subprocess.TimeoutExpired:
+                print("Server did not terminate gracefully, killing.")
+                process.kill()
+
+            # Print any remaining output for debugging
+            remaining_output = process.stdout.read()
+            if remaining_output:
+                print("\n--- Remaining Server Output ---")
+                print(remaining_output)
+