ComfyUI/tests-unit/app_test/frontend_manager_test.py

import argparse
from pathlib import Path

import pytest
from requests.exceptions import HTTPError
from unittest.mock import patch, mock_open

from app.frontend_management import (
    FrontendManager,
    FrontEndProvider,
    Release,
)
from comfy.cli_args import DEFAULT_VERSION_STRING


@pytest.fixture
def mock_releases():
    return [
        Release(
            id=1,
            tag_name="1.0.0",
            name="Release 1.0.0",
            prerelease=False,
            created_at="2022-01-01T00:00:00Z",
            published_at="2022-01-01T00:00:00Z",
            body="Release notes for 1.0.0",
            assets=[{"name": "dist.zip", "url": "https://example.com/dist.zip"}],
        ),
        Release(
            id=2,
            tag_name="2.0.0",
            name="Release 2.0.0",
            prerelease=False,
            created_at="2022-02-01T00:00:00Z",
            published_at="2022-02-01T00:00:00Z",
            body="Release notes for 2.0.0",
            assets=[{"name": "dist.zip", "url": "https://example.com/dist.zip"}],
        ),
    ]


@pytest.fixture
def mock_provider(mock_releases):
    provider = FrontEndProvider(
        owner="test-owner",
        repo="test-repo",
    )
    provider.all_releases = mock_releases
    provider.latest_release = mock_releases[1]
    FrontendManager.PROVIDERS = [provider]
    return provider


@pytest.fixture(autouse=True)
def clear_cache():
    import utils.install_util
    import app.frontend_management

    utils.install_util.PACKAGE_VERSIONS = {}
    app.frontend_management.COMFY_PACKAGE_VERSIONS = []


def test_get_release(mock_provider, mock_releases):
    version = "1.0.0"
    release = mock_provider.get_release(version)
    assert release == mock_releases[0]


def test_get_release_latest(mock_provider, mock_releases):
    version = "latest"
    release = mock_provider.get_release(version)
    assert release == mock_releases[1]


def test_get_release_invalid_version(mock_provider):
    version = "invalid"
    with pytest.raises(ValueError):
        mock_provider.get_release(version)


def test_init_frontend_default():
    version_string = DEFAULT_VERSION_STRING
    frontend_path = FrontendManager.init_frontend(version_string)
    assert frontend_path == FrontendManager.default_frontend_path()


def test_init_frontend_invalid_version():
    version_string = "test-owner/test-repo@1.100.99"
    with pytest.raises(HTTPError):
        FrontendManager.init_frontend_unsafe(version_string)


def test_init_frontend_invalid_provider():
    version_string = "invalid/invalid@latest"
    with pytest.raises(HTTPError):
        FrontendManager.init_frontend_unsafe(version_string)


@pytest.fixture
def mock_os_functions():
    with (
        patch("app.frontend_management.os.makedirs") as mock_makedirs,
        patch("app.frontend_management.os.listdir") as mock_listdir,
        patch("app.frontend_management.os.rmdir") as mock_rmdir,
    ):
        mock_listdir.return_value = []  # Simulate empty directory
        yield mock_makedirs, mock_listdir, mock_rmdir


@pytest.fixture
def mock_download():
    with patch("app.frontend_management.download_release_asset_zip") as mock:
        mock.side_effect = Exception("Download failed")  # Simulate download failure
        yield mock


def test_finally_block(mock_os_functions, mock_download, mock_provider):
    # Arrange
    mock_makedirs, mock_listdir, mock_rmdir = mock_os_functions
    version_string = "test-owner/test-repo@1.0.0"

    # Act & Assert
    with pytest.raises(Exception):
        FrontendManager.init_frontend_unsafe(version_string, mock_provider)

    # Assert
    mock_makedirs.assert_called_once()
    mock_download.assert_called_once()
    mock_listdir.assert_called_once()
    mock_rmdir.assert_called_once()


def test_parse_version_string():
    version_string = "owner/repo@1.0.0"
    repo_owner, repo_name, version = FrontendManager.parse_version_string(
        version_string
    )
    assert repo_owner == "owner"
    assert repo_name == "repo"
    assert version == "1.0.0"


def test_parse_version_string_invalid():
    version_string = "invalid"
    with pytest.raises(argparse.ArgumentTypeError):
        FrontendManager.parse_version_string(version_string)


def test_init_frontend_default_with_mocks():
    # Arrange
    version_string = DEFAULT_VERSION_STRING

    # Act
    with (
        patch("app.frontend_management.check_comfy_packages_versions") as mock_check,
        patch.object(
            FrontendManager, "default_frontend_path", return_value="/mocked/path"
        ),
    ):
        frontend_path = FrontendManager.init_frontend(version_string)

    # Assert
    assert frontend_path == "/mocked/path"
    mock_check.assert_called_once()


def test_init_frontend_fallback_on_error():
    # Arrange
    version_string = "test-owner/test-repo@1.0.0"

    # Act
    with (
        patch.object(
            FrontendManager, "init_frontend_unsafe", side_effect=Exception("Test error")
        ),
        patch("app.frontend_management.check_comfy_packages_versions") as mock_check,
        patch.object(
            FrontendManager, "default_frontend_path", return_value="/default/path"
        ),
    ):
        frontend_path = FrontendManager.init_frontend(version_string)

    # Assert
    assert frontend_path == "/default/path"
    mock_check.assert_called_once()


def test_get_frontend_version():
    # Arrange
    expected_version = "1.25.0"
    mock_requirements_content = """torch
torchsde
comfyui-frontend-package==1.25.0
other-package==1.0.0
numpy"""

    # Act
    with patch("builtins.open", mock_open(read_data=mock_requirements_content)):
        version = FrontendManager.get_required_frontend_version()

    # Assert
    assert version == expected_version


def test_get_frontend_version_invalid_semver():
    # Arrange
    mock_requirements_content = """torch
torchsde
comfyui-frontend-package==1.29.3.75
other-package==1.0.0
numpy"""

    # Act
    with patch("builtins.open", mock_open(read_data=mock_requirements_content)):
        version = FrontendManager.get_required_frontend_version()

    # Assert
    assert version is None


def test_get_templates_version():
    # Arrange
    expected_version = "0.1.41"
    mock_requirements_content = """torch
torchsde
comfyui-frontend-package==1.25.0
comfyui-workflow-templates==0.1.41
other-package==1.0.0
numpy"""

    # Act
    with patch("builtins.open", mock_open(read_data=mock_requirements_content)):
        version = FrontendManager.get_required_templates_version()

    # Assert
    assert version == expected_version


def test_get_templates_version_not_found():
    # Arrange
    mock_requirements_content = """torch
torchsde
comfyui-frontend-package==1.25.0
other-package==1.0.0
numpy"""

    # Act
    with patch("builtins.open", mock_open(read_data=mock_requirements_content)):
        version = FrontendManager.get_required_templates_version()

    # Assert
    assert version is None


def test_get_templates_version_invalid_semver():
    # Arrange
    mock_requirements_content = """torch
torchsde
comfyui-workflow-templates==1.0.0.beta
other-package==1.0.0
numpy"""

    # Act
    with patch("builtins.open", mock_open(read_data=mock_requirements_content)):
        version = FrontendManager.get_required_templates_version()

    # Assert
    assert version is None


def test_get_installed_templates_version():
    # Arrange
    expected_version = "0.1.40"

    # Act
    with patch("app.frontend_management.version", return_value=expected_version):
        version = FrontendManager.get_installed_templates_version()

    # Assert
    assert version == expected_version


def test_get_installed_templates_version_not_installed():
    # Act
    with patch(
        "app.frontend_management.version", side_effect=Exception("Package not found")
    ):
        version = FrontendManager.get_installed_templates_version()

    # Assert
    assert version is None


# ---------------------------------------------------------------------------
# Auto-managed @latest / @prerelease cleanup (CORE-285)
# ---------------------------------------------------------------------------


@pytest.fixture
def custom_frontends_root(tmp_path, monkeypatch):
    """Point ``FrontendManager.CUSTOM_FRONTENDS_ROOT`` at a fresh tmp dir."""
    root = tmp_path / "web_custom_versions"
    root.mkdir()
    monkeypatch.setattr(FrontendManager, "CUSTOM_FRONTENDS_ROOT", str(root))
    return root


def _make_version_dir(root, owner, repo, version):
    """Create ``<root>/<owner>_<repo>/<version>/index.html`` and return path."""
    folder = root / f"{owner}_{repo}" / version
    folder.mkdir(parents=True, exist_ok=True)
    (folder / "index.html").write_text("<html></html>")
    return folder


def test_auto_managed_metadata_roundtrip(custom_frontends_root):
    FrontendManager._write_auto_managed_versions("o", "r", ["1.0.0", "1.1.0", "1.0.0"])
    assert FrontendManager._read_auto_managed_versions("o", "r") == ["1.0.0", "1.1.0"]


def test_read_auto_managed_versions_missing(custom_frontends_root):
    assert FrontendManager._read_auto_managed_versions("o", "r") == []


def test_read_auto_managed_versions_corrupt(custom_frontends_root):
    provider_dir = custom_frontends_root / "o_r"
    provider_dir.mkdir()
    (provider_dir / FrontendManager.AUTO_MANAGED_METADATA_FILENAME).write_text(
        "not json"
    )
    assert FrontendManager._read_auto_managed_versions("o", "r") == []


@pytest.mark.parametrize("payload", ["[]", "null", '"oops"', "42"])
def test_read_auto_managed_versions_non_dict_root(custom_frontends_root, payload):
    """Valid JSON whose root isn't a dict must not raise — bug pointed out in
    Oracle review: ``data.get(...)`` would throw on non-dict roots."""
    provider_dir = custom_frontends_root / "o_r"
    provider_dir.mkdir()
    (provider_dir / FrontendManager.AUTO_MANAGED_METADATA_FILENAME).write_text(payload)
    assert FrontendManager._read_auto_managed_versions("o", "r") == []


def test_read_auto_managed_versions_drops_unsafe_names(custom_frontends_root):
    """Tampered metadata containing path-traversal-y names must be ignored
    at read time so cleanup never even sees them."""
    provider_dir = custom_frontends_root / "o_r"
    provider_dir.mkdir()
    (provider_dir / FrontendManager.AUTO_MANAGED_METADATA_FILENAME).write_text(
        '{"auto_managed": ["1.0.0", "../escape", "/etc", "..", ".", "ok-1.2"]}'
    )
    assert FrontendManager._read_auto_managed_versions("o", "r") == ["1.0.0", "ok-1.2"]


def test_write_auto_managed_versions_drops_unsafe_names(custom_frontends_root):
    """Even if a caller passes a tainted list, the file we persist must
    only contain safe names."""
    FrontendManager._write_auto_managed_versions(
        "o", "r", ["1.0.0", "../escape", "/etc/passwd", "..", "."]
    )
    assert FrontendManager._read_auto_managed_versions("o", "r") == ["1.0.0"]


def test_prune_refuses_to_delete_outside_provider_dir(
    custom_frontends_root, monkeypatch
):
    """Defense in depth: even if a hostile name slips past the read/write
    filters somehow, ``_prune_auto_managed_versions`` must refuse to rmtree
    anything outside the provider directory."""
    # Set up a sibling directory that must NOT be touched.
    sibling = custom_frontends_root / "do-not-touch"
    sibling.mkdir()
    (sibling / "marker").write_text("keep me")

    provider_dir = custom_frontends_root / "o_r"
    provider_dir.mkdir()
    _make_version_dir(custom_frontends_root, "o", "r", "1.0.0")

    # Bypass the read filter to simulate a corrupt-but-parseable list slipping
    # through (e.g. older code wrote it before this hardening landed).
    monkeypatch.setattr(
        FrontendManager,
        "_read_auto_managed_versions",
        classmethod(lambda cls, owner, repo: ["1.0.0", "../do-not-touch"]),
    )

    FrontendManager._prune_auto_managed_versions("o", "r", keep_version="2.0.0")

    # Sibling untouched.
    assert sibling.exists()
    assert (sibling / "marker").read_text() == "keep me"
    # In-bounds folder still gets cleaned.
    assert not (provider_dir / "1.0.0").exists()


def test_prune_auto_managed_versions_removes_stale_and_keeps_pinned(
    custom_frontends_root,
):
    # Two versions previously fetched via @latest, plus an explicitly pinned one.
    _make_version_dir(custom_frontends_root, "o", "r", "1.0.0")
    _make_version_dir(custom_frontends_root, "o", "r", "1.1.0")
    pinned = _make_version_dir(custom_frontends_root, "o", "r", "1.2.0")
    FrontendManager._write_auto_managed_versions("o", "r", ["1.0.0", "1.1.0"])

    # User runs @latest again and it resolves to 1.1.0 — older auto-managed
    # 1.0.0 should be deleted, pinned 1.2.0 should remain untouched.
    FrontendManager._prune_auto_managed_versions("o", "r", keep_version="1.1.0")

    provider_dir = custom_frontends_root / "o_r"
    assert not (provider_dir / "1.0.0").exists()
    assert (provider_dir / "1.1.0").exists()
    assert pinned.exists()
    assert FrontendManager._read_auto_managed_versions("o", "r") == ["1.1.0"]


def test_untrack_auto_managed_version_does_not_delete_folder(custom_frontends_root):
    version_dir = _make_version_dir(custom_frontends_root, "o", "r", "1.0.0")
    FrontendManager._write_auto_managed_versions("o", "r", ["1.0.0", "1.1.0"])

    FrontendManager._untrack_auto_managed_version("o", "r", "1.0.0")

    assert version_dir.exists()
    assert FrontendManager._read_auto_managed_versions("o", "r") == ["1.1.0"]


def test_init_frontend_latest_prunes_previous_auto_managed_versions(
    custom_frontends_root, mock_provider, mock_releases
):
    # Pre-existing folders: 1.0.0 was previously downloaded via @latest, 1.1.5
    # was explicitly pinned by the user. Now @latest resolves to 2.0.0.
    _make_version_dir(custom_frontends_root, "test-owner", "test-repo", "1.0.0")
    pinned = _make_version_dir(
        custom_frontends_root, "test-owner", "test-repo", "1.1.5"
    )
    FrontendManager._write_auto_managed_versions(
        "test-owner", "test-repo", ["1.0.0"]
    )

    # Stub out the actual download so we just create the destination dir.
    def fake_download(release, destination_path):
        Path(destination_path).mkdir(parents=True, exist_ok=True)
        (Path(destination_path) / "index.html").write_text("<html></html>")

    with patch(
        "app.frontend_management.download_release_asset_zip",
        side_effect=fake_download,
    ):
        result = FrontendManager.init_frontend_unsafe(
            "test-owner/test-repo@latest", mock_provider
        )

    provider_dir = custom_frontends_root / "test-owner_test-repo"
    # 2.0.0 was downloaded and tracked.
    assert Path(result) == provider_dir / "2.0.0"
    assert (provider_dir / "2.0.0").exists()
    assert FrontendManager._read_auto_managed_versions(
        "test-owner", "test-repo"
    ) == ["2.0.0"]
    # Old auto-managed 1.0.0 was pruned.
    assert not (provider_dir / "1.0.0").exists()
    # Pinned 1.1.5 was left alone.
    assert pinned.exists()


def test_init_frontend_explicit_version_promotes_out_of_auto_managed(
    custom_frontends_root, mock_provider
):
    # 1.0.0 was previously downloaded via @latest.
    _make_version_dir(custom_frontends_root, "test-owner", "test-repo", "1.0.0")
    FrontendManager._write_auto_managed_versions(
        "test-owner", "test-repo", ["1.0.0"]
    )

    # User now explicitly pins it. The `v`-prefixed early-return path runs.
    result = FrontendManager.init_frontend_unsafe(
        "test-owner/test-repo@v1.0.0", mock_provider
    )

    provider_dir = custom_frontends_root / "test-owner_test-repo"
    assert Path(result) == provider_dir / "1.0.0"
    # It should no longer be tracked as auto-managed, so a future @latest run
    # won't sweep it away.
    assert FrontendManager._read_auto_managed_versions(
        "test-owner", "test-repo"
    ) == []
    # The folder is still on disk.
    assert (provider_dir / "1.0.0").exists()


def test_init_frontend_explicit_version_no_v_prefix_promotes_out_of_auto_managed(
    custom_frontends_root, mock_provider
):
    # 1.0.0 was previously downloaded via @latest, and is also already on
    # disk so the GitHub-resolution path is a no-op for download.
    _make_version_dir(custom_frontends_root, "test-owner", "test-repo", "1.0.0")
    FrontendManager._write_auto_managed_versions(
        "test-owner", "test-repo", ["1.0.0"]
    )

    # No `v` prefix → goes through the GitHub release lookup path. The folder
    # already exists, so download is skipped.
    with patch(
        "app.frontend_management.download_release_asset_zip"
    ) as mock_download_zip:
        FrontendManager.init_frontend_unsafe(
            "test-owner/test-repo@1.0.0", mock_provider
        )
        mock_download_zip.assert_not_called()

    # It should be promoted out of auto-managed even when the folder is reused.
    assert FrontendManager._read_auto_managed_versions(
        "test-owner", "test-repo"
    ) == []