mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-02-19 22:44:35 +00:00
8ff4d38ad1
This change solves the basename collision bug by using UNIQUE(file_path) on the unified asset_references table. Key changes: Database: - Migration 0005 merges asset_cache_states and asset_infos into asset_references - AssetReference now contains: cache state fields (file_path, mtime_ns, needs_verify, is_missing, enrichment_level) plus info fields (name, owner_id, preview_id, etc.) - AssetReferenceMeta replaces AssetInfoMeta - AssetReferenceTag replaces AssetInfoTag - UNIQUE constraint on file_path prevents duplicate entries for same file Code: - New unified query module: asset_reference.py (replaces asset_info.py, cache_state.py) - Updated scanner, seeder, and services to use AssetReference - Updated API routes to use reference_id instead of asset_info_id Tests: - All 175 unit tests updated and passing - Integration tests require server environment (not run here) Amp-Thread-ID: https://ampcode.com/threads/T-019c4fe8-9dcb-75ce-bea8-ea786343a581 Co-authored-by: Amp <amp@ampcode.com>
254 lines
7.7 KiB
Python
254 lines
7.7 KiB
Python
"""Tests for asset enrichment (mime_type and hash population)."""
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
from sqlalchemy.orm import Session
|
|
|
|
from app.assets.database.models import Asset, AssetReference
|
|
from app.assets.scanner import (
|
|
ENRICHMENT_HASHED,
|
|
ENRICHMENT_METADATA,
|
|
ENRICHMENT_STUB,
|
|
enrich_asset,
|
|
)
|
|
|
|
|
|
def _create_stub_asset(
|
|
session: Session,
|
|
file_path: str,
|
|
asset_id: str = "test-asset-id",
|
|
reference_id: str = "test-ref-id",
|
|
name: str | None = None,
|
|
) -> tuple[Asset, AssetReference]:
|
|
"""Create a stub asset with reference for testing enrichment."""
|
|
asset = Asset(
|
|
id=asset_id,
|
|
hash=None,
|
|
size_bytes=100,
|
|
mime_type=None,
|
|
)
|
|
session.add(asset)
|
|
session.flush()
|
|
|
|
ref = AssetReference(
|
|
id=reference_id,
|
|
asset_id=asset_id,
|
|
name=name or f"test-asset-{asset_id}",
|
|
owner_id="system",
|
|
file_path=file_path,
|
|
mtime_ns=1234567890000000000,
|
|
enrichment_level=ENRICHMENT_STUB,
|
|
)
|
|
session.add(ref)
|
|
session.flush()
|
|
|
|
return asset, ref
|
|
|
|
|
|
class TestEnrichAsset:
|
|
def test_extracts_mime_type_and_updates_asset(
|
|
self, db_engine, temp_dir: Path, session: Session
|
|
):
|
|
"""Verify mime_type is written to the Asset table during enrichment."""
|
|
file_path = temp_dir / "model.safetensors"
|
|
file_path.write_bytes(b"\x00" * 100)
|
|
|
|
asset, ref = _create_stub_asset(
|
|
session, str(file_path), "asset-1", "ref-1"
|
|
)
|
|
session.commit()
|
|
|
|
with patch("app.assets.scanner.create_session") as mock_cs:
|
|
from contextlib import contextmanager
|
|
|
|
@contextmanager
|
|
def _create_session():
|
|
with Session(db_engine) as sess:
|
|
yield sess
|
|
|
|
mock_cs.side_effect = _create_session
|
|
|
|
new_level = enrich_asset(
|
|
file_path=str(file_path),
|
|
reference_id=ref.id,
|
|
asset_id=asset.id,
|
|
extract_metadata=True,
|
|
compute_hash=False,
|
|
)
|
|
|
|
assert new_level == ENRICHMENT_METADATA
|
|
|
|
session.expire_all()
|
|
updated_asset = session.get(Asset, "asset-1")
|
|
assert updated_asset is not None
|
|
assert updated_asset.mime_type == "application/safetensors"
|
|
|
|
def test_computes_hash_and_updates_asset(
|
|
self, db_engine, temp_dir: Path, session: Session
|
|
):
|
|
"""Verify hash is written to the Asset table during enrichment."""
|
|
file_path = temp_dir / "data.bin"
|
|
file_path.write_bytes(b"test content for hashing")
|
|
|
|
asset, ref = _create_stub_asset(
|
|
session, str(file_path), "asset-2", "ref-2"
|
|
)
|
|
session.commit()
|
|
|
|
with patch("app.assets.scanner.create_session") as mock_cs:
|
|
from contextlib import contextmanager
|
|
|
|
@contextmanager
|
|
def _create_session():
|
|
with Session(db_engine) as sess:
|
|
yield sess
|
|
|
|
mock_cs.side_effect = _create_session
|
|
|
|
new_level = enrich_asset(
|
|
file_path=str(file_path),
|
|
reference_id=ref.id,
|
|
asset_id=asset.id,
|
|
extract_metadata=True,
|
|
compute_hash=True,
|
|
)
|
|
|
|
assert new_level == ENRICHMENT_HASHED
|
|
|
|
session.expire_all()
|
|
updated_asset = session.get(Asset, "asset-2")
|
|
assert updated_asset is not None
|
|
assert updated_asset.hash is not None
|
|
assert updated_asset.hash.startswith("blake3:")
|
|
|
|
def test_enrichment_updates_both_mime_and_hash(
|
|
self, db_engine, temp_dir: Path, session: Session
|
|
):
|
|
"""Verify both mime_type and hash are set when full enrichment runs."""
|
|
file_path = temp_dir / "model.safetensors"
|
|
file_path.write_bytes(b"\x00" * 50)
|
|
|
|
asset, ref = _create_stub_asset(
|
|
session, str(file_path), "asset-3", "ref-3"
|
|
)
|
|
session.commit()
|
|
|
|
with patch("app.assets.scanner.create_session") as mock_cs:
|
|
from contextlib import contextmanager
|
|
|
|
@contextmanager
|
|
def _create_session():
|
|
with Session(db_engine) as sess:
|
|
yield sess
|
|
|
|
mock_cs.side_effect = _create_session
|
|
|
|
enrich_asset(
|
|
file_path=str(file_path),
|
|
reference_id=ref.id,
|
|
asset_id=asset.id,
|
|
extract_metadata=True,
|
|
compute_hash=True,
|
|
)
|
|
|
|
session.expire_all()
|
|
updated_asset = session.get(Asset, "asset-3")
|
|
assert updated_asset is not None
|
|
assert updated_asset.mime_type == "application/safetensors"
|
|
assert updated_asset.hash is not None
|
|
assert updated_asset.hash.startswith("blake3:")
|
|
|
|
def test_missing_file_returns_stub_level(
|
|
self, db_engine, temp_dir: Path, session: Session
|
|
):
|
|
"""Verify missing files don't cause errors and return STUB level."""
|
|
file_path = temp_dir / "nonexistent.bin"
|
|
|
|
asset, ref = _create_stub_asset(
|
|
session, str(file_path), "asset-4", "ref-4"
|
|
)
|
|
session.commit()
|
|
|
|
with patch("app.assets.scanner.create_session") as mock_cs:
|
|
from contextlib import contextmanager
|
|
|
|
@contextmanager
|
|
def _create_session():
|
|
with Session(db_engine) as sess:
|
|
yield sess
|
|
|
|
mock_cs.side_effect = _create_session
|
|
|
|
new_level = enrich_asset(
|
|
file_path=str(file_path),
|
|
reference_id=ref.id,
|
|
asset_id=asset.id,
|
|
extract_metadata=True,
|
|
compute_hash=True,
|
|
)
|
|
|
|
assert new_level == ENRICHMENT_STUB
|
|
|
|
session.expire_all()
|
|
updated_asset = session.get(Asset, "asset-4")
|
|
assert updated_asset.mime_type is None
|
|
assert updated_asset.hash is None
|
|
|
|
def test_duplicate_hash_merges_into_existing_asset(
|
|
self, db_engine, temp_dir: Path, session: Session
|
|
):
|
|
"""Verify duplicate files merge into existing asset instead of failing."""
|
|
file_path_1 = temp_dir / "file1.bin"
|
|
file_path_2 = temp_dir / "file2.bin"
|
|
content = b"identical content"
|
|
file_path_1.write_bytes(content)
|
|
file_path_2.write_bytes(content)
|
|
|
|
asset1, ref1 = _create_stub_asset(
|
|
session, str(file_path_1), "asset-dup-1", "ref-dup-1"
|
|
)
|
|
asset2, ref2 = _create_stub_asset(
|
|
session, str(file_path_2), "asset-dup-2", "ref-dup-2"
|
|
)
|
|
session.commit()
|
|
|
|
with patch("app.assets.scanner.create_session") as mock_cs:
|
|
from contextlib import contextmanager
|
|
|
|
@contextmanager
|
|
def _create_session():
|
|
with Session(db_engine) as sess:
|
|
yield sess
|
|
|
|
mock_cs.side_effect = _create_session
|
|
|
|
enrich_asset(
|
|
file_path=str(file_path_1),
|
|
reference_id=ref1.id,
|
|
asset_id=asset1.id,
|
|
extract_metadata=True,
|
|
compute_hash=True,
|
|
)
|
|
|
|
enrich_asset(
|
|
file_path=str(file_path_2),
|
|
reference_id=ref2.id,
|
|
asset_id=asset2.id,
|
|
extract_metadata=True,
|
|
compute_hash=True,
|
|
)
|
|
|
|
session.expire_all()
|
|
|
|
updated_asset1 = session.get(Asset, "asset-dup-1")
|
|
assert updated_asset1 is not None
|
|
assert updated_asset1.hash is not None
|
|
|
|
updated_asset2 = session.get(Asset, "asset-dup-2")
|
|
assert updated_asset2 is None
|
|
|
|
updated_ref2 = session.get(AssetReference, "ref-dup-2")
|
|
assert updated_ref2 is not None
|
|
assert updated_ref2.asset_id == "asset-dup-1"
|