mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-03-05 17:27:42 +08:00
4d4c2cedd3
- Fix missing import for compute_filename_for_reference in ingest.py - Apply code review fixes across routes, queries, scanner, seeder, hashing, ingest, path_utils, main, and server - Update and add tests for sync references and seeder Amp-Thread-ID: https://ampcode.com/threads/T-019cb61a-ed54-738c-a05f-9b5242e513f3 Co-authored-by: Amp <amp@ampcode.com>
351 lines
11 KiB
Python
351 lines
11 KiB
Python
"""Tests for sync_references_with_filesystem in scanner.py."""
|
|
|
|
import os
|
|
import tempfile
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
from sqlalchemy import create_engine
|
|
from sqlalchemy.orm import Session
|
|
|
|
from app.assets.database.models import (
|
|
Asset,
|
|
AssetReference,
|
|
AssetReferenceTag,
|
|
Base,
|
|
Tag,
|
|
)
|
|
from app.assets.scanner import sync_references_with_filesystem
|
|
from app.assets.services.file_utils import get_mtime_ns
|
|
|
|
|
|
@pytest.fixture
|
|
def db_engine():
|
|
engine = create_engine("sqlite:///:memory:")
|
|
Base.metadata.create_all(engine)
|
|
return engine
|
|
|
|
|
|
@pytest.fixture
|
|
def session(db_engine):
|
|
with Session(db_engine) as sess:
|
|
yield sess
|
|
|
|
|
|
@pytest.fixture
|
|
def temp_dir():
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
yield Path(tmpdir)
|
|
|
|
|
|
def _create_file(temp_dir: Path, name: str, content: bytes = b"\x00" * 100) -> str:
|
|
"""Create a file and return its absolute path (no symlink resolution)."""
|
|
p = temp_dir / name
|
|
p.parent.mkdir(parents=True, exist_ok=True)
|
|
p.write_bytes(content)
|
|
return os.path.abspath(str(p))
|
|
|
|
|
|
def _stat_mtime_ns(path: str) -> int:
|
|
return get_mtime_ns(os.stat(path, follow_symlinks=True))
|
|
|
|
|
|
def _make_asset(
|
|
session: Session,
|
|
asset_id: str,
|
|
file_path: str,
|
|
ref_id: str,
|
|
*,
|
|
asset_hash: str | None = None,
|
|
size_bytes: int = 100,
|
|
mtime_ns: int | None = None,
|
|
needs_verify: bool = False,
|
|
is_missing: bool = False,
|
|
) -> tuple[Asset, AssetReference]:
|
|
"""Insert an Asset + AssetReference and flush."""
|
|
asset = session.get(Asset, asset_id)
|
|
if asset is None:
|
|
asset = Asset(id=asset_id, hash=asset_hash, size_bytes=size_bytes)
|
|
session.add(asset)
|
|
session.flush()
|
|
|
|
ref = AssetReference(
|
|
id=ref_id,
|
|
asset_id=asset_id,
|
|
name=f"test-{ref_id}",
|
|
owner_id="system",
|
|
file_path=file_path,
|
|
mtime_ns=mtime_ns,
|
|
needs_verify=needs_verify,
|
|
is_missing=is_missing,
|
|
)
|
|
session.add(ref)
|
|
session.flush()
|
|
return asset, ref
|
|
|
|
|
|
def _ensure_missing_tag(session: Session):
|
|
"""Ensure the 'missing' tag exists."""
|
|
if not session.get(Tag, "missing"):
|
|
session.add(Tag(name="missing", tag_type="system"))
|
|
session.flush()
|
|
|
|
|
|
class _VerifyCase:
|
|
def __init__(self, id, stat_unchanged, needs_verify_before, expect_needs_verify):
|
|
self.id = id
|
|
self.stat_unchanged = stat_unchanged
|
|
self.needs_verify_before = needs_verify_before
|
|
self.expect_needs_verify = expect_needs_verify
|
|
|
|
|
|
VERIFY_CASES = [
|
|
_VerifyCase(
|
|
id="unchanged_clears_verify",
|
|
stat_unchanged=True,
|
|
needs_verify_before=True,
|
|
expect_needs_verify=False,
|
|
),
|
|
_VerifyCase(
|
|
id="unchanged_keeps_clear",
|
|
stat_unchanged=True,
|
|
needs_verify_before=False,
|
|
expect_needs_verify=False,
|
|
),
|
|
_VerifyCase(
|
|
id="changed_sets_verify",
|
|
stat_unchanged=False,
|
|
needs_verify_before=False,
|
|
expect_needs_verify=True,
|
|
),
|
|
_VerifyCase(
|
|
id="changed_keeps_verify",
|
|
stat_unchanged=False,
|
|
needs_verify_before=True,
|
|
expect_needs_verify=True,
|
|
),
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize("case", VERIFY_CASES, ids=lambda c: c.id)
|
|
def test_needs_verify_toggling(session, temp_dir, case):
|
|
"""needs_verify is set/cleared based on mtime+size match."""
|
|
fp = _create_file(temp_dir, "model.bin")
|
|
real_mtime = _stat_mtime_ns(fp)
|
|
|
|
mtime_for_db = real_mtime if case.stat_unchanged else real_mtime + 1
|
|
_make_asset(
|
|
session, "a1", fp, "r1",
|
|
asset_hash="blake3:abc",
|
|
mtime_ns=mtime_for_db,
|
|
needs_verify=case.needs_verify_before,
|
|
)
|
|
session.commit()
|
|
|
|
with patch("app.assets.scanner.get_prefixes_for_root", return_value=[str(temp_dir)]):
|
|
sync_references_with_filesystem(session, "models")
|
|
session.commit()
|
|
|
|
session.expire_all()
|
|
ref = session.get(AssetReference, "r1")
|
|
assert ref.needs_verify is case.expect_needs_verify
|
|
|
|
|
|
class _MissingCase:
|
|
def __init__(self, id, file_exists, expect_is_missing):
|
|
self.id = id
|
|
self.file_exists = file_exists
|
|
self.expect_is_missing = expect_is_missing
|
|
|
|
|
|
MISSING_CASES = [
|
|
_MissingCase(id="existing_file_not_missing", file_exists=True, expect_is_missing=False),
|
|
_MissingCase(id="missing_file_marked_missing", file_exists=False, expect_is_missing=True),
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize("case", MISSING_CASES, ids=lambda c: c.id)
|
|
def test_is_missing_flag(session, temp_dir, case):
|
|
"""is_missing reflects whether the file exists on disk."""
|
|
if case.file_exists:
|
|
fp = _create_file(temp_dir, "model.bin")
|
|
mtime = _stat_mtime_ns(fp)
|
|
else:
|
|
fp = str(temp_dir / "gone.bin")
|
|
mtime = 999
|
|
|
|
_make_asset(session, "a1", fp, "r1", asset_hash="blake3:abc", mtime_ns=mtime)
|
|
session.commit()
|
|
|
|
with patch("app.assets.scanner.get_prefixes_for_root", return_value=[str(temp_dir)]):
|
|
sync_references_with_filesystem(session, "models")
|
|
session.commit()
|
|
|
|
session.expire_all()
|
|
ref = session.get(AssetReference, "r1")
|
|
assert ref.is_missing is case.expect_is_missing
|
|
|
|
|
|
def test_seed_asset_all_missing_deletes_asset(session, temp_dir):
|
|
"""Seed asset with all refs missing gets deleted entirely."""
|
|
fp = str(temp_dir / "gone.bin")
|
|
_make_asset(session, "seed1", fp, "r1", asset_hash=None, mtime_ns=999)
|
|
session.commit()
|
|
|
|
with patch("app.assets.scanner.get_prefixes_for_root", return_value=[str(temp_dir)]):
|
|
sync_references_with_filesystem(session, "models")
|
|
session.commit()
|
|
|
|
assert session.get(Asset, "seed1") is None
|
|
assert session.get(AssetReference, "r1") is None
|
|
|
|
|
|
def test_seed_asset_some_exist_returns_survivors(session, temp_dir):
|
|
"""Seed asset with at least one existing ref survives and is returned."""
|
|
fp = _create_file(temp_dir, "model.bin")
|
|
mtime = _stat_mtime_ns(fp)
|
|
_make_asset(session, "seed1", fp, "r1", asset_hash=None, mtime_ns=mtime)
|
|
session.commit()
|
|
|
|
with patch("app.assets.scanner.get_prefixes_for_root", return_value=[str(temp_dir)]):
|
|
survivors = sync_references_with_filesystem(
|
|
session, "models", collect_existing_paths=True,
|
|
)
|
|
session.commit()
|
|
|
|
assert session.get(Asset, "seed1") is not None
|
|
assert os.path.abspath(fp) in survivors
|
|
|
|
|
|
def test_hashed_asset_prunes_missing_refs_when_one_is_ok(session, temp_dir):
|
|
"""Hashed asset with one stat-unchanged ref deletes missing refs."""
|
|
fp_ok = _create_file(temp_dir, "good.bin")
|
|
fp_gone = str(temp_dir / "gone.bin")
|
|
mtime = _stat_mtime_ns(fp_ok)
|
|
|
|
_make_asset(session, "h1", fp_ok, "r_ok", asset_hash="blake3:aaa", mtime_ns=mtime)
|
|
# Second ref on same asset, file missing
|
|
ref_gone = AssetReference(
|
|
id="r_gone", asset_id="h1", name="gone",
|
|
owner_id="system", file_path=fp_gone, mtime_ns=999,
|
|
)
|
|
session.add(ref_gone)
|
|
session.commit()
|
|
|
|
with patch("app.assets.scanner.get_prefixes_for_root", return_value=[str(temp_dir)]):
|
|
sync_references_with_filesystem(session, "models")
|
|
session.commit()
|
|
|
|
session.expire_all()
|
|
assert session.get(AssetReference, "r_ok") is not None
|
|
assert session.get(AssetReference, "r_gone") is None
|
|
|
|
|
|
def test_hashed_asset_all_missing_keeps_refs(session, temp_dir):
|
|
"""Hashed asset with all refs missing keeps refs (no pruning)."""
|
|
fp = str(temp_dir / "gone.bin")
|
|
_make_asset(session, "h1", fp, "r1", asset_hash="blake3:aaa", mtime_ns=999)
|
|
session.commit()
|
|
|
|
with patch("app.assets.scanner.get_prefixes_for_root", return_value=[str(temp_dir)]):
|
|
sync_references_with_filesystem(session, "models")
|
|
session.commit()
|
|
|
|
session.expire_all()
|
|
assert session.get(AssetReference, "r1") is not None
|
|
ref = session.get(AssetReference, "r1")
|
|
assert ref.is_missing is True
|
|
|
|
|
|
def test_missing_tag_added_when_all_refs_gone(session, temp_dir):
|
|
"""Missing tag is added to hashed asset when all refs are missing."""
|
|
_ensure_missing_tag(session)
|
|
fp = str(temp_dir / "gone.bin")
|
|
_make_asset(session, "h1", fp, "r1", asset_hash="blake3:aaa", mtime_ns=999)
|
|
session.commit()
|
|
|
|
with patch("app.assets.scanner.get_prefixes_for_root", return_value=[str(temp_dir)]):
|
|
sync_references_with_filesystem(
|
|
session, "models", update_missing_tags=True,
|
|
)
|
|
session.commit()
|
|
|
|
session.expire_all()
|
|
tag_link = session.get(AssetReferenceTag, ("r1", "missing"))
|
|
assert tag_link is not None
|
|
|
|
|
|
def test_missing_tag_removed_when_ref_ok(session, temp_dir):
|
|
"""Missing tag is removed from hashed asset when a ref is stat-unchanged."""
|
|
_ensure_missing_tag(session)
|
|
fp = _create_file(temp_dir, "model.bin")
|
|
mtime = _stat_mtime_ns(fp)
|
|
_make_asset(session, "h1", fp, "r1", asset_hash="blake3:aaa", mtime_ns=mtime)
|
|
# Pre-add a stale missing tag
|
|
session.add(AssetReferenceTag(
|
|
asset_reference_id="r1", tag_name="missing", origin="automatic",
|
|
))
|
|
session.commit()
|
|
|
|
with patch("app.assets.scanner.get_prefixes_for_root", return_value=[str(temp_dir)]):
|
|
sync_references_with_filesystem(
|
|
session, "models", update_missing_tags=True,
|
|
)
|
|
session.commit()
|
|
|
|
session.expire_all()
|
|
tag_link = session.get(AssetReferenceTag, ("r1", "missing"))
|
|
assert tag_link is None
|
|
|
|
|
|
def test_missing_tags_not_touched_when_flag_false(session, temp_dir):
|
|
"""Missing tags are not modified when update_missing_tags=False."""
|
|
_ensure_missing_tag(session)
|
|
fp = str(temp_dir / "gone.bin")
|
|
_make_asset(session, "h1", fp, "r1", asset_hash="blake3:aaa", mtime_ns=999)
|
|
session.commit()
|
|
|
|
with patch("app.assets.scanner.get_prefixes_for_root", return_value=[str(temp_dir)]):
|
|
sync_references_with_filesystem(
|
|
session, "models", update_missing_tags=False,
|
|
)
|
|
session.commit()
|
|
|
|
tag_link = session.get(AssetReferenceTag, ("r1", "missing"))
|
|
assert tag_link is None # tag was never added
|
|
|
|
|
|
def test_returns_none_when_collect_false(session, temp_dir):
|
|
fp = _create_file(temp_dir, "model.bin")
|
|
mtime = _stat_mtime_ns(fp)
|
|
_make_asset(session, "a1", fp, "r1", asset_hash="blake3:abc", mtime_ns=mtime)
|
|
session.commit()
|
|
|
|
with patch("app.assets.scanner.get_prefixes_for_root", return_value=[str(temp_dir)]):
|
|
result = sync_references_with_filesystem(
|
|
session, "models", collect_existing_paths=False,
|
|
)
|
|
|
|
assert result is None
|
|
|
|
|
|
def test_returns_empty_set_for_no_prefixes(session):
|
|
with patch("app.assets.scanner.get_prefixes_for_root", return_value=[]):
|
|
result = sync_references_with_filesystem(
|
|
session, "models", collect_existing_paths=True,
|
|
)
|
|
|
|
assert result == set()
|
|
|
|
|
|
def test_no_references_is_noop(session, temp_dir):
|
|
"""No crash and no side effects when there are no references."""
|
|
with patch("app.assets.scanner.get_prefixes_for_root", return_value=[str(temp_dir)]):
|
|
survivors = sync_references_with_filesystem(
|
|
session, "models", collect_existing_paths=True,
|
|
)
|
|
session.commit()
|
|
|
|
assert survivors == set()
|