mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-04-25 00:49:13 +00:00
refactor(assets): split queries.py into modular query modules
Split the ~1000 line app/assets/database/queries.py into focused modules: - queries/asset.py - Asset entity queries (asset_exists_by_hash, get_asset_by_hash) - queries/asset_info.py - AssetInfo queries (~15 functions) - queries/cache_state.py - AssetCacheState queries (list_cache_states_by_asset_id, pick_best_live_path, prune_orphaned_assets, fast_db_consistency_pass) - queries/tags.py - Tag queries (8 functions including ensure_tags_exist, add/remove tag functions, list_tags_with_usage) - queries/__init__.py - Re-exports all public functions for backward compatibility Also adds comprehensive unit tests using in-memory SQLite: - tests-unit/assets_test/queries/conftest.py - Session fixture - tests-unit/assets_test/queries/test_asset.py - 5 tests - tests-unit/assets_test/queries/test_asset_info.py - 23 tests - tests-unit/assets_test/queries/test_cache_state.py - 8 tests - tests-unit/assets_test/queries/test_metadata.py - 12 tests for _apply_metadata_filter - tests-unit/assets_test/queries/test_tags.py - 23 tests All 71 unit tests pass. Existing integration tests unaffected. Amp-Thread-ID: https://ampcode.com/threads/T-019c24bb-475b-7442-9ff9-8288edea3345 Co-authored-by: Amp <amp@ampcode.com>
This commit is contained in:
14
tests-unit/assets_test/queries/conftest.py
Normal file
14
tests-unit/assets_test/queries/conftest.py
Normal file
@@ -0,0 +1,14 @@
|
||||
import pytest
|
||||
from sqlalchemy import create_engine
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from app.assets.database.models import Base
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def session():
|
||||
"""In-memory SQLite session for fast unit tests."""
|
||||
engine = create_engine("sqlite:///:memory:")
|
||||
Base.metadata.create_all(engine)
|
||||
with Session(engine) as sess:
|
||||
yield sess
|
||||
39
tests-unit/assets_test/queries/test_asset.py
Normal file
39
tests-unit/assets_test/queries/test_asset.py
Normal file
@@ -0,0 +1,39 @@
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from app.assets.database.models import Asset
|
||||
from app.assets.database.queries import asset_exists_by_hash, get_asset_by_hash
|
||||
|
||||
|
||||
class TestAssetExistsByHash:
|
||||
def test_returns_false_for_nonexistent(self, session: Session):
|
||||
assert asset_exists_by_hash(session, asset_hash="nonexistent") is False
|
||||
|
||||
def test_returns_true_for_existing(self, session: Session):
|
||||
asset = Asset(hash="blake3:abc123", size_bytes=100)
|
||||
session.add(asset)
|
||||
session.commit()
|
||||
|
||||
assert asset_exists_by_hash(session, asset_hash="blake3:abc123") is True
|
||||
|
||||
def test_does_not_match_null_hash(self, session: Session):
|
||||
asset = Asset(hash=None, size_bytes=100)
|
||||
session.add(asset)
|
||||
session.commit()
|
||||
|
||||
assert asset_exists_by_hash(session, asset_hash="") is False
|
||||
|
||||
|
||||
class TestGetAssetByHash:
|
||||
def test_returns_none_for_nonexistent(self, session: Session):
|
||||
assert get_asset_by_hash(session, asset_hash="nonexistent") is None
|
||||
|
||||
def test_returns_asset_for_existing(self, session: Session):
|
||||
asset = Asset(hash="blake3:def456", size_bytes=200, mime_type="image/png")
|
||||
session.add(asset)
|
||||
session.commit()
|
||||
|
||||
result = get_asset_by_hash(session, asset_hash="blake3:def456")
|
||||
assert result is not None
|
||||
assert result.id == asset.id
|
||||
assert result.size_bytes == 200
|
||||
assert result.mime_type == "image/png"
|
||||
268
tests-unit/assets_test/queries/test_asset_info.py
Normal file
268
tests-unit/assets_test/queries/test_asset_info.py
Normal file
@@ -0,0 +1,268 @@
|
||||
import pytest
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from app.assets.database.models import Asset, AssetInfo, AssetInfoMeta, AssetInfoTag, Tag
|
||||
from app.assets.database.queries import (
|
||||
asset_info_exists_for_asset_id,
|
||||
get_asset_info_by_id,
|
||||
list_asset_infos_page,
|
||||
fetch_asset_info_asset_and_tags,
|
||||
fetch_asset_info_and_asset,
|
||||
touch_asset_info_by_id,
|
||||
delete_asset_info_by_id,
|
||||
set_asset_info_preview,
|
||||
ensure_tags_exist,
|
||||
add_tags_to_asset_info,
|
||||
)
|
||||
from app.assets.helpers import utcnow
|
||||
|
||||
|
||||
def _make_asset(session: Session, hash_val: str | None = None, size: int = 1024) -> Asset:
|
||||
asset = Asset(hash=hash_val, size_bytes=size, mime_type="application/octet-stream")
|
||||
session.add(asset)
|
||||
session.flush()
|
||||
return asset
|
||||
|
||||
|
||||
def _make_asset_info(
|
||||
session: Session,
|
||||
asset: Asset,
|
||||
name: str = "test",
|
||||
owner_id: str = "",
|
||||
) -> AssetInfo:
|
||||
now = utcnow()
|
||||
info = AssetInfo(
|
||||
owner_id=owner_id,
|
||||
name=name,
|
||||
asset_id=asset.id,
|
||||
created_at=now,
|
||||
updated_at=now,
|
||||
last_access_time=now,
|
||||
)
|
||||
session.add(info)
|
||||
session.flush()
|
||||
return info
|
||||
|
||||
|
||||
class TestAssetInfoExistsForAssetId:
|
||||
def test_returns_false_when_no_info(self, session: Session):
|
||||
asset = _make_asset(session, "hash1")
|
||||
assert asset_info_exists_for_asset_id(session, asset_id=asset.id) is False
|
||||
|
||||
def test_returns_true_when_info_exists(self, session: Session):
|
||||
asset = _make_asset(session, "hash1")
|
||||
_make_asset_info(session, asset)
|
||||
assert asset_info_exists_for_asset_id(session, asset_id=asset.id) is True
|
||||
|
||||
|
||||
class TestGetAssetInfoById:
|
||||
def test_returns_none_for_nonexistent(self, session: Session):
|
||||
assert get_asset_info_by_id(session, asset_info_id="nonexistent") is None
|
||||
|
||||
def test_returns_info(self, session: Session):
|
||||
asset = _make_asset(session, "hash1")
|
||||
info = _make_asset_info(session, asset, name="myfile.txt")
|
||||
|
||||
result = get_asset_info_by_id(session, asset_info_id=info.id)
|
||||
assert result is not None
|
||||
assert result.name == "myfile.txt"
|
||||
|
||||
|
||||
class TestListAssetInfosPage:
|
||||
def test_empty_db(self, session: Session):
|
||||
infos, tag_map, total = list_asset_infos_page(session)
|
||||
assert infos == []
|
||||
assert tag_map == {}
|
||||
assert total == 0
|
||||
|
||||
def test_returns_infos_with_tags(self, session: Session):
|
||||
asset = _make_asset(session, "hash1")
|
||||
info = _make_asset_info(session, asset, name="test.bin")
|
||||
ensure_tags_exist(session, ["alpha", "beta"])
|
||||
add_tags_to_asset_info(session, asset_info_id=info.id, tags=["alpha", "beta"])
|
||||
session.commit()
|
||||
|
||||
infos, tag_map, total = list_asset_infos_page(session)
|
||||
assert len(infos) == 1
|
||||
assert infos[0].id == info.id
|
||||
assert set(tag_map[info.id]) == {"alpha", "beta"}
|
||||
assert total == 1
|
||||
|
||||
def test_name_contains_filter(self, session: Session):
|
||||
asset = _make_asset(session, "hash1")
|
||||
_make_asset_info(session, asset, name="model_v1.safetensors")
|
||||
_make_asset_info(session, asset, name="config.json")
|
||||
session.commit()
|
||||
|
||||
infos, _, total = list_asset_infos_page(session, name_contains="model")
|
||||
assert total == 1
|
||||
assert infos[0].name == "model_v1.safetensors"
|
||||
|
||||
def test_owner_visibility(self, session: Session):
|
||||
asset = _make_asset(session, "hash1")
|
||||
_make_asset_info(session, asset, name="public", owner_id="")
|
||||
_make_asset_info(session, asset, name="private", owner_id="user1")
|
||||
session.commit()
|
||||
|
||||
# Empty owner sees only public
|
||||
infos, _, total = list_asset_infos_page(session, owner_id="")
|
||||
assert total == 1
|
||||
assert infos[0].name == "public"
|
||||
|
||||
# Owner sees both
|
||||
infos, _, total = list_asset_infos_page(session, owner_id="user1")
|
||||
assert total == 2
|
||||
|
||||
def test_include_tags_filter(self, session: Session):
|
||||
asset = _make_asset(session, "hash1")
|
||||
info1 = _make_asset_info(session, asset, name="tagged")
|
||||
info2 = _make_asset_info(session, asset, name="untagged")
|
||||
ensure_tags_exist(session, ["wanted"])
|
||||
add_tags_to_asset_info(session, asset_info_id=info1.id, tags=["wanted"])
|
||||
session.commit()
|
||||
|
||||
infos, _, total = list_asset_infos_page(session, include_tags=["wanted"])
|
||||
assert total == 1
|
||||
assert infos[0].name == "tagged"
|
||||
|
||||
def test_exclude_tags_filter(self, session: Session):
|
||||
asset = _make_asset(session, "hash1")
|
||||
info1 = _make_asset_info(session, asset, name="keep")
|
||||
info2 = _make_asset_info(session, asset, name="exclude")
|
||||
ensure_tags_exist(session, ["bad"])
|
||||
add_tags_to_asset_info(session, asset_info_id=info2.id, tags=["bad"])
|
||||
session.commit()
|
||||
|
||||
infos, _, total = list_asset_infos_page(session, exclude_tags=["bad"])
|
||||
assert total == 1
|
||||
assert infos[0].name == "keep"
|
||||
|
||||
def test_sorting(self, session: Session):
|
||||
asset = _make_asset(session, "hash1", size=100)
|
||||
asset2 = _make_asset(session, "hash2", size=500)
|
||||
_make_asset_info(session, asset, name="small")
|
||||
_make_asset_info(session, asset2, name="large")
|
||||
session.commit()
|
||||
|
||||
infos, _, _ = list_asset_infos_page(session, sort="size", order="desc")
|
||||
assert infos[0].name == "large"
|
||||
|
||||
infos, _, _ = list_asset_infos_page(session, sort="name", order="asc")
|
||||
assert infos[0].name == "large"
|
||||
|
||||
|
||||
class TestFetchAssetInfoAssetAndTags:
|
||||
def test_returns_none_for_nonexistent(self, session: Session):
|
||||
result = fetch_asset_info_asset_and_tags(session, "nonexistent")
|
||||
assert result is None
|
||||
|
||||
def test_returns_tuple(self, session: Session):
|
||||
asset = _make_asset(session, "hash1")
|
||||
info = _make_asset_info(session, asset, name="test.bin")
|
||||
ensure_tags_exist(session, ["tag1"])
|
||||
add_tags_to_asset_info(session, asset_info_id=info.id, tags=["tag1"])
|
||||
session.commit()
|
||||
|
||||
result = fetch_asset_info_asset_and_tags(session, info.id)
|
||||
assert result is not None
|
||||
ret_info, ret_asset, ret_tags = result
|
||||
assert ret_info.id == info.id
|
||||
assert ret_asset.id == asset.id
|
||||
assert ret_tags == ["tag1"]
|
||||
|
||||
|
||||
class TestFetchAssetInfoAndAsset:
|
||||
def test_returns_none_for_nonexistent(self, session: Session):
|
||||
result = fetch_asset_info_and_asset(session, asset_info_id="nonexistent")
|
||||
assert result is None
|
||||
|
||||
def test_returns_tuple(self, session: Session):
|
||||
asset = _make_asset(session, "hash1")
|
||||
info = _make_asset_info(session, asset)
|
||||
session.commit()
|
||||
|
||||
result = fetch_asset_info_and_asset(session, asset_info_id=info.id)
|
||||
assert result is not None
|
||||
ret_info, ret_asset = result
|
||||
assert ret_info.id == info.id
|
||||
assert ret_asset.id == asset.id
|
||||
|
||||
|
||||
class TestTouchAssetInfoById:
|
||||
def test_updates_last_access_time(self, session: Session):
|
||||
asset = _make_asset(session, "hash1")
|
||||
info = _make_asset_info(session, asset)
|
||||
original_time = info.last_access_time
|
||||
session.commit()
|
||||
|
||||
import time
|
||||
time.sleep(0.01)
|
||||
|
||||
touch_asset_info_by_id(session, asset_info_id=info.id)
|
||||
session.commit()
|
||||
|
||||
session.refresh(info)
|
||||
assert info.last_access_time > original_time
|
||||
|
||||
|
||||
class TestDeleteAssetInfoById:
|
||||
def test_deletes_existing(self, session: Session):
|
||||
asset = _make_asset(session, "hash1")
|
||||
info = _make_asset_info(session, asset)
|
||||
session.commit()
|
||||
|
||||
result = delete_asset_info_by_id(session, asset_info_id=info.id, owner_id="")
|
||||
assert result is True
|
||||
assert get_asset_info_by_id(session, asset_info_id=info.id) is None
|
||||
|
||||
def test_returns_false_for_nonexistent(self, session: Session):
|
||||
result = delete_asset_info_by_id(session, asset_info_id="nonexistent", owner_id="")
|
||||
assert result is False
|
||||
|
||||
def test_respects_owner_visibility(self, session: Session):
|
||||
asset = _make_asset(session, "hash1")
|
||||
info = _make_asset_info(session, asset, owner_id="user1")
|
||||
session.commit()
|
||||
|
||||
result = delete_asset_info_by_id(session, asset_info_id=info.id, owner_id="user2")
|
||||
assert result is False
|
||||
assert get_asset_info_by_id(session, asset_info_id=info.id) is not None
|
||||
|
||||
|
||||
class TestSetAssetInfoPreview:
|
||||
def test_sets_preview(self, session: Session):
|
||||
asset = _make_asset(session, "hash1")
|
||||
preview_asset = _make_asset(session, "preview_hash")
|
||||
info = _make_asset_info(session, asset)
|
||||
session.commit()
|
||||
|
||||
set_asset_info_preview(session, asset_info_id=info.id, preview_asset_id=preview_asset.id)
|
||||
session.commit()
|
||||
|
||||
session.refresh(info)
|
||||
assert info.preview_id == preview_asset.id
|
||||
|
||||
def test_clears_preview(self, session: Session):
|
||||
asset = _make_asset(session, "hash1")
|
||||
preview_asset = _make_asset(session, "preview_hash")
|
||||
info = _make_asset_info(session, asset)
|
||||
info.preview_id = preview_asset.id
|
||||
session.commit()
|
||||
|
||||
set_asset_info_preview(session, asset_info_id=info.id, preview_asset_id=None)
|
||||
session.commit()
|
||||
|
||||
session.refresh(info)
|
||||
assert info.preview_id is None
|
||||
|
||||
def test_raises_for_nonexistent_info(self, session: Session):
|
||||
with pytest.raises(ValueError, match="not found"):
|
||||
set_asset_info_preview(session, asset_info_id="nonexistent", preview_asset_id=None)
|
||||
|
||||
def test_raises_for_nonexistent_preview(self, session: Session):
|
||||
asset = _make_asset(session, "hash1")
|
||||
info = _make_asset_info(session, asset)
|
||||
session.commit()
|
||||
|
||||
with pytest.raises(ValueError, match="Preview Asset"):
|
||||
set_asset_info_preview(session, asset_info_id=info.id, preview_asset_id="nonexistent")
|
||||
128
tests-unit/assets_test/queries/test_cache_state.py
Normal file
128
tests-unit/assets_test/queries/test_cache_state.py
Normal file
@@ -0,0 +1,128 @@
|
||||
"""Tests for cache_state query functions."""
|
||||
import os
|
||||
import tempfile
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from app.assets.database.models import Asset, AssetCacheState, AssetInfo
|
||||
from app.assets.database.queries import (
|
||||
list_cache_states_by_asset_id,
|
||||
pick_best_live_path,
|
||||
)
|
||||
from app.assets.helpers import utcnow
|
||||
|
||||
|
||||
def _make_asset(session: Session, hash_val: str | None = None, size: int = 1024) -> Asset:
|
||||
asset = Asset(hash=hash_val, size_bytes=size)
|
||||
session.add(asset)
|
||||
session.flush()
|
||||
return asset
|
||||
|
||||
|
||||
def _make_cache_state(
|
||||
session: Session,
|
||||
asset: Asset,
|
||||
file_path: str,
|
||||
mtime_ns: int | None = None,
|
||||
needs_verify: bool = False,
|
||||
) -> AssetCacheState:
|
||||
state = AssetCacheState(
|
||||
asset_id=asset.id,
|
||||
file_path=file_path,
|
||||
mtime_ns=mtime_ns,
|
||||
needs_verify=needs_verify,
|
||||
)
|
||||
session.add(state)
|
||||
session.flush()
|
||||
return state
|
||||
|
||||
|
||||
class TestListCacheStatesByAssetId:
|
||||
def test_returns_empty_for_no_states(self, session: Session):
|
||||
asset = _make_asset(session, "hash1")
|
||||
states = list_cache_states_by_asset_id(session, asset_id=asset.id)
|
||||
assert list(states) == []
|
||||
|
||||
def test_returns_states_for_asset(self, session: Session):
|
||||
asset = _make_asset(session, "hash1")
|
||||
_make_cache_state(session, asset, "/path/a.bin")
|
||||
_make_cache_state(session, asset, "/path/b.bin")
|
||||
session.commit()
|
||||
|
||||
states = list_cache_states_by_asset_id(session, asset_id=asset.id)
|
||||
paths = [s.file_path for s in states]
|
||||
assert set(paths) == {"/path/a.bin", "/path/b.bin"}
|
||||
|
||||
def test_does_not_return_other_assets_states(self, session: Session):
|
||||
asset1 = _make_asset(session, "hash1")
|
||||
asset2 = _make_asset(session, "hash2")
|
||||
_make_cache_state(session, asset1, "/path/asset1.bin")
|
||||
_make_cache_state(session, asset2, "/path/asset2.bin")
|
||||
session.commit()
|
||||
|
||||
states = list_cache_states_by_asset_id(session, asset_id=asset1.id)
|
||||
paths = [s.file_path for s in states]
|
||||
assert paths == ["/path/asset1.bin"]
|
||||
|
||||
|
||||
class TestPickBestLivePath:
|
||||
def test_returns_empty_for_empty_list(self):
|
||||
result = pick_best_live_path([])
|
||||
assert result == ""
|
||||
|
||||
def test_returns_empty_when_no_files_exist(self, session: Session):
|
||||
asset = _make_asset(session, "hash1")
|
||||
state = _make_cache_state(session, asset, "/nonexistent/path.bin")
|
||||
session.commit()
|
||||
|
||||
result = pick_best_live_path([state])
|
||||
assert result == ""
|
||||
|
||||
def test_prefers_verified_path(self, session: Session, tmp_path):
|
||||
"""needs_verify=False should be preferred."""
|
||||
asset = _make_asset(session, "hash1")
|
||||
|
||||
verified_file = tmp_path / "verified.bin"
|
||||
verified_file.write_bytes(b"data")
|
||||
|
||||
unverified_file = tmp_path / "unverified.bin"
|
||||
unverified_file.write_bytes(b"data")
|
||||
|
||||
state_verified = _make_cache_state(
|
||||
session, asset, str(verified_file), needs_verify=False
|
||||
)
|
||||
state_unverified = _make_cache_state(
|
||||
session, asset, str(unverified_file), needs_verify=True
|
||||
)
|
||||
session.commit()
|
||||
|
||||
states = [state_unverified, state_verified]
|
||||
result = pick_best_live_path(states)
|
||||
assert result == str(verified_file)
|
||||
|
||||
def test_falls_back_to_existing_unverified(self, session: Session, tmp_path):
|
||||
"""If all states need verification, return first existing path."""
|
||||
asset = _make_asset(session, "hash1")
|
||||
|
||||
existing_file = tmp_path / "exists.bin"
|
||||
existing_file.write_bytes(b"data")
|
||||
|
||||
state = _make_cache_state(session, asset, str(existing_file), needs_verify=True)
|
||||
session.commit()
|
||||
|
||||
result = pick_best_live_path([state])
|
||||
assert result == str(existing_file)
|
||||
|
||||
|
||||
class TestPickBestLivePathWithMocking:
|
||||
def test_handles_missing_file_path_attr(self):
|
||||
"""Gracefully handle states with None file_path."""
|
||||
|
||||
class MockState:
|
||||
file_path = None
|
||||
needs_verify = False
|
||||
|
||||
result = pick_best_live_path([MockState()])
|
||||
assert result == ""
|
||||
180
tests-unit/assets_test/queries/test_metadata.py
Normal file
180
tests-unit/assets_test/queries/test_metadata.py
Normal file
@@ -0,0 +1,180 @@
|
||||
"""Tests for metadata filtering logic in asset_info queries."""
|
||||
import pytest
|
||||
from decimal import Decimal
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from app.assets.database.models import Asset, AssetInfo, AssetInfoMeta
|
||||
from app.assets.database.queries import list_asset_infos_page
|
||||
from app.assets.helpers import utcnow, project_kv
|
||||
|
||||
|
||||
def _make_asset(session: Session, hash_val: str) -> Asset:
|
||||
asset = Asset(hash=hash_val, size_bytes=1024)
|
||||
session.add(asset)
|
||||
session.flush()
|
||||
return asset
|
||||
|
||||
|
||||
def _make_asset_info(
|
||||
session: Session,
|
||||
asset: Asset,
|
||||
name: str,
|
||||
metadata: dict | None = None,
|
||||
) -> AssetInfo:
|
||||
now = utcnow()
|
||||
info = AssetInfo(
|
||||
owner_id="",
|
||||
name=name,
|
||||
asset_id=asset.id,
|
||||
user_metadata=metadata,
|
||||
created_at=now,
|
||||
updated_at=now,
|
||||
last_access_time=now,
|
||||
)
|
||||
session.add(info)
|
||||
session.flush()
|
||||
|
||||
if metadata:
|
||||
for key, val in metadata.items():
|
||||
for row in project_kv(key, val):
|
||||
meta_row = AssetInfoMeta(
|
||||
asset_info_id=info.id,
|
||||
key=row["key"],
|
||||
ordinal=row.get("ordinal", 0),
|
||||
val_str=row.get("val_str"),
|
||||
val_num=row.get("val_num"),
|
||||
val_bool=row.get("val_bool"),
|
||||
val_json=row.get("val_json"),
|
||||
)
|
||||
session.add(meta_row)
|
||||
session.flush()
|
||||
|
||||
return info
|
||||
|
||||
|
||||
class TestMetadataFilterString:
|
||||
def test_filter_by_string_value(self, session: Session):
|
||||
asset = _make_asset(session, "hash1")
|
||||
_make_asset_info(session, asset, "match", {"category": "models"})
|
||||
_make_asset_info(session, asset, "nomatch", {"category": "images"})
|
||||
session.commit()
|
||||
|
||||
infos, _, total = list_asset_infos_page(session, metadata_filter={"category": "models"})
|
||||
assert total == 1
|
||||
assert infos[0].name == "match"
|
||||
|
||||
def test_filter_by_string_no_match(self, session: Session):
|
||||
asset = _make_asset(session, "hash1")
|
||||
_make_asset_info(session, asset, "item", {"category": "models"})
|
||||
session.commit()
|
||||
|
||||
infos, _, total = list_asset_infos_page(session, metadata_filter={"category": "other"})
|
||||
assert total == 0
|
||||
|
||||
|
||||
class TestMetadataFilterNumeric:
|
||||
def test_filter_by_int_value(self, session: Session):
|
||||
asset = _make_asset(session, "hash1")
|
||||
_make_asset_info(session, asset, "epoch5", {"epoch": 5})
|
||||
_make_asset_info(session, asset, "epoch10", {"epoch": 10})
|
||||
session.commit()
|
||||
|
||||
infos, _, total = list_asset_infos_page(session, metadata_filter={"epoch": 5})
|
||||
assert total == 1
|
||||
assert infos[0].name == "epoch5"
|
||||
|
||||
def test_filter_by_float_value(self, session: Session):
|
||||
asset = _make_asset(session, "hash1")
|
||||
_make_asset_info(session, asset, "high", {"score": 0.95})
|
||||
_make_asset_info(session, asset, "low", {"score": 0.5})
|
||||
session.commit()
|
||||
|
||||
infos, _, total = list_asset_infos_page(session, metadata_filter={"score": 0.95})
|
||||
assert total == 1
|
||||
assert infos[0].name == "high"
|
||||
|
||||
|
||||
class TestMetadataFilterBoolean:
|
||||
def test_filter_by_true(self, session: Session):
|
||||
asset = _make_asset(session, "hash1")
|
||||
_make_asset_info(session, asset, "active", {"enabled": True})
|
||||
_make_asset_info(session, asset, "inactive", {"enabled": False})
|
||||
session.commit()
|
||||
|
||||
infos, _, total = list_asset_infos_page(session, metadata_filter={"enabled": True})
|
||||
assert total == 1
|
||||
assert infos[0].name == "active"
|
||||
|
||||
def test_filter_by_false(self, session: Session):
|
||||
asset = _make_asset(session, "hash1")
|
||||
_make_asset_info(session, asset, "active", {"enabled": True})
|
||||
_make_asset_info(session, asset, "inactive", {"enabled": False})
|
||||
session.commit()
|
||||
|
||||
infos, _, total = list_asset_infos_page(session, metadata_filter={"enabled": False})
|
||||
assert total == 1
|
||||
assert infos[0].name == "inactive"
|
||||
|
||||
|
||||
class TestMetadataFilterNull:
|
||||
def test_filter_by_null_matches_missing_key(self, session: Session):
|
||||
asset = _make_asset(session, "hash1")
|
||||
_make_asset_info(session, asset, "has_key", {"optional": "value"})
|
||||
_make_asset_info(session, asset, "missing_key", {})
|
||||
session.commit()
|
||||
|
||||
infos, _, total = list_asset_infos_page(session, metadata_filter={"optional": None})
|
||||
assert total == 1
|
||||
assert infos[0].name == "missing_key"
|
||||
|
||||
def test_filter_by_null_matches_explicit_null(self, session: Session):
|
||||
asset = _make_asset(session, "hash1")
|
||||
_make_asset_info(session, asset, "explicit_null", {"nullable": None})
|
||||
_make_asset_info(session, asset, "has_value", {"nullable": "present"})
|
||||
session.commit()
|
||||
|
||||
infos, _, total = list_asset_infos_page(session, metadata_filter={"nullable": None})
|
||||
assert total == 1
|
||||
assert infos[0].name == "explicit_null"
|
||||
|
||||
|
||||
class TestMetadataFilterList:
|
||||
def test_filter_by_list_or(self, session: Session):
|
||||
"""List values should match ANY of the values (OR)."""
|
||||
asset = _make_asset(session, "hash1")
|
||||
_make_asset_info(session, asset, "cat_a", {"category": "a"})
|
||||
_make_asset_info(session, asset, "cat_b", {"category": "b"})
|
||||
_make_asset_info(session, asset, "cat_c", {"category": "c"})
|
||||
session.commit()
|
||||
|
||||
infos, _, total = list_asset_infos_page(session, metadata_filter={"category": ["a", "b"]})
|
||||
assert total == 2
|
||||
names = {i.name for i in infos}
|
||||
assert names == {"cat_a", "cat_b"}
|
||||
|
||||
|
||||
class TestMetadataFilterMultipleKeys:
|
||||
def test_multiple_keys_and(self, session: Session):
|
||||
"""Multiple keys should ALL match (AND)."""
|
||||
asset = _make_asset(session, "hash1")
|
||||
_make_asset_info(session, asset, "match", {"type": "model", "version": 2})
|
||||
_make_asset_info(session, asset, "wrong_type", {"type": "config", "version": 2})
|
||||
_make_asset_info(session, asset, "wrong_version", {"type": "model", "version": 1})
|
||||
session.commit()
|
||||
|
||||
infos, _, total = list_asset_infos_page(
|
||||
session, metadata_filter={"type": "model", "version": 2}
|
||||
)
|
||||
assert total == 1
|
||||
assert infos[0].name == "match"
|
||||
|
||||
|
||||
class TestMetadataFilterEmptyDict:
|
||||
def test_empty_filter_returns_all(self, session: Session):
|
||||
asset = _make_asset(session, "hash1")
|
||||
_make_asset_info(session, asset, "a", {"key": "val"})
|
||||
_make_asset_info(session, asset, "b", {})
|
||||
session.commit()
|
||||
|
||||
infos, _, total = list_asset_infos_page(session, metadata_filter={})
|
||||
assert total == 2
|
||||
297
tests-unit/assets_test/queries/test_tags.py
Normal file
297
tests-unit/assets_test/queries/test_tags.py
Normal file
@@ -0,0 +1,297 @@
|
||||
import pytest
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from app.assets.database.models import Asset, AssetInfo, AssetInfoTag, Tag
|
||||
from app.assets.database.queries import (
|
||||
ensure_tags_exist,
|
||||
get_asset_tags,
|
||||
set_asset_info_tags,
|
||||
add_tags_to_asset_info,
|
||||
remove_tags_from_asset_info,
|
||||
add_missing_tag_for_asset_id,
|
||||
remove_missing_tag_for_asset_id,
|
||||
list_tags_with_usage,
|
||||
)
|
||||
from app.assets.helpers import utcnow
|
||||
|
||||
|
||||
def _make_asset(session: Session, hash_val: str | None = None) -> Asset:
|
||||
asset = Asset(hash=hash_val, size_bytes=1024)
|
||||
session.add(asset)
|
||||
session.flush()
|
||||
return asset
|
||||
|
||||
|
||||
def _make_asset_info(session: Session, asset: Asset, name: str = "test", owner_id: str = "") -> AssetInfo:
|
||||
now = utcnow()
|
||||
info = AssetInfo(
|
||||
owner_id=owner_id,
|
||||
name=name,
|
||||
asset_id=asset.id,
|
||||
created_at=now,
|
||||
updated_at=now,
|
||||
last_access_time=now,
|
||||
)
|
||||
session.add(info)
|
||||
session.flush()
|
||||
return info
|
||||
|
||||
|
||||
class TestEnsureTagsExist:
|
||||
def test_creates_new_tags(self, session: Session):
|
||||
ensure_tags_exist(session, ["alpha", "beta"], tag_type="user")
|
||||
session.commit()
|
||||
|
||||
tags = session.query(Tag).all()
|
||||
assert {t.name for t in tags} == {"alpha", "beta"}
|
||||
|
||||
def test_is_idempotent(self, session: Session):
|
||||
ensure_tags_exist(session, ["alpha"], tag_type="user")
|
||||
ensure_tags_exist(session, ["alpha"], tag_type="user")
|
||||
session.commit()
|
||||
|
||||
assert session.query(Tag).count() == 1
|
||||
|
||||
def test_normalizes_tags(self, session: Session):
|
||||
ensure_tags_exist(session, [" ALPHA ", "Beta", "alpha"])
|
||||
session.commit()
|
||||
|
||||
tags = session.query(Tag).all()
|
||||
assert {t.name for t in tags} == {"alpha", "beta"}
|
||||
|
||||
def test_empty_list_is_noop(self, session: Session):
|
||||
ensure_tags_exist(session, [])
|
||||
session.commit()
|
||||
assert session.query(Tag).count() == 0
|
||||
|
||||
def test_tag_type_is_set(self, session: Session):
|
||||
ensure_tags_exist(session, ["system-tag"], tag_type="system")
|
||||
session.commit()
|
||||
|
||||
tag = session.query(Tag).filter_by(name="system-tag").one()
|
||||
assert tag.tag_type == "system"
|
||||
|
||||
|
||||
class TestGetAssetTags:
|
||||
def test_returns_empty_for_no_tags(self, session: Session):
|
||||
asset = _make_asset(session, "hash1")
|
||||
info = _make_asset_info(session, asset)
|
||||
|
||||
tags = get_asset_tags(session, asset_info_id=info.id)
|
||||
assert tags == []
|
||||
|
||||
def test_returns_tags_for_asset(self, session: Session):
|
||||
asset = _make_asset(session, "hash1")
|
||||
info = _make_asset_info(session, asset)
|
||||
|
||||
ensure_tags_exist(session, ["tag1", "tag2"])
|
||||
session.add_all([
|
||||
AssetInfoTag(asset_info_id=info.id, tag_name="tag1", origin="manual", added_at=utcnow()),
|
||||
AssetInfoTag(asset_info_id=info.id, tag_name="tag2", origin="manual", added_at=utcnow()),
|
||||
])
|
||||
session.flush()
|
||||
|
||||
tags = get_asset_tags(session, asset_info_id=info.id)
|
||||
assert set(tags) == {"tag1", "tag2"}
|
||||
|
||||
|
||||
class TestSetAssetInfoTags:
|
||||
def test_adds_new_tags(self, session: Session):
|
||||
asset = _make_asset(session, "hash1")
|
||||
info = _make_asset_info(session, asset)
|
||||
|
||||
result = set_asset_info_tags(session, asset_info_id=info.id, tags=["a", "b"])
|
||||
session.commit()
|
||||
|
||||
assert set(result["added"]) == {"a", "b"}
|
||||
assert result["removed"] == []
|
||||
assert set(result["total"]) == {"a", "b"}
|
||||
|
||||
def test_removes_old_tags(self, session: Session):
|
||||
asset = _make_asset(session, "hash1")
|
||||
info = _make_asset_info(session, asset)
|
||||
|
||||
set_asset_info_tags(session, asset_info_id=info.id, tags=["a", "b", "c"])
|
||||
result = set_asset_info_tags(session, asset_info_id=info.id, tags=["a"])
|
||||
session.commit()
|
||||
|
||||
assert result["added"] == []
|
||||
assert set(result["removed"]) == {"b", "c"}
|
||||
assert result["total"] == ["a"]
|
||||
|
||||
def test_replaces_tags(self, session: Session):
|
||||
asset = _make_asset(session, "hash1")
|
||||
info = _make_asset_info(session, asset)
|
||||
|
||||
set_asset_info_tags(session, asset_info_id=info.id, tags=["a", "b"])
|
||||
result = set_asset_info_tags(session, asset_info_id=info.id, tags=["b", "c"])
|
||||
session.commit()
|
||||
|
||||
assert result["added"] == ["c"]
|
||||
assert result["removed"] == ["a"]
|
||||
assert set(result["total"]) == {"b", "c"}
|
||||
|
||||
|
||||
class TestAddTagsToAssetInfo:
|
||||
def test_adds_tags(self, session: Session):
|
||||
asset = _make_asset(session, "hash1")
|
||||
info = _make_asset_info(session, asset)
|
||||
|
||||
result = add_tags_to_asset_info(session, asset_info_id=info.id, tags=["x", "y"])
|
||||
session.commit()
|
||||
|
||||
assert set(result["added"]) == {"x", "y"}
|
||||
assert result["already_present"] == []
|
||||
|
||||
def test_reports_already_present(self, session: Session):
|
||||
asset = _make_asset(session, "hash1")
|
||||
info = _make_asset_info(session, asset)
|
||||
|
||||
add_tags_to_asset_info(session, asset_info_id=info.id, tags=["x"])
|
||||
result = add_tags_to_asset_info(session, asset_info_id=info.id, tags=["x", "y"])
|
||||
session.commit()
|
||||
|
||||
assert result["added"] == ["y"]
|
||||
assert result["already_present"] == ["x"]
|
||||
|
||||
def test_raises_for_missing_asset_info(self, session: Session):
|
||||
with pytest.raises(ValueError, match="not found"):
|
||||
add_tags_to_asset_info(session, asset_info_id="nonexistent", tags=["x"])
|
||||
|
||||
|
||||
class TestRemoveTagsFromAssetInfo:
|
||||
def test_removes_tags(self, session: Session):
|
||||
asset = _make_asset(session, "hash1")
|
||||
info = _make_asset_info(session, asset)
|
||||
|
||||
add_tags_to_asset_info(session, asset_info_id=info.id, tags=["a", "b", "c"])
|
||||
result = remove_tags_from_asset_info(session, asset_info_id=info.id, tags=["a", "b"])
|
||||
session.commit()
|
||||
|
||||
assert set(result["removed"]) == {"a", "b"}
|
||||
assert result["not_present"] == []
|
||||
assert result["total_tags"] == ["c"]
|
||||
|
||||
def test_reports_not_present(self, session: Session):
|
||||
asset = _make_asset(session, "hash1")
|
||||
info = _make_asset_info(session, asset)
|
||||
|
||||
add_tags_to_asset_info(session, asset_info_id=info.id, tags=["a"])
|
||||
result = remove_tags_from_asset_info(session, asset_info_id=info.id, tags=["a", "x"])
|
||||
session.commit()
|
||||
|
||||
assert result["removed"] == ["a"]
|
||||
assert result["not_present"] == ["x"]
|
||||
|
||||
def test_raises_for_missing_asset_info(self, session: Session):
|
||||
with pytest.raises(ValueError, match="not found"):
|
||||
remove_tags_from_asset_info(session, asset_info_id="nonexistent", tags=["x"])
|
||||
|
||||
|
||||
class TestMissingTagFunctions:
|
||||
def test_add_missing_tag_for_asset_id(self, session: Session):
|
||||
asset = _make_asset(session, "hash1")
|
||||
info = _make_asset_info(session, asset)
|
||||
ensure_tags_exist(session, ["missing"], tag_type="system")
|
||||
|
||||
add_missing_tag_for_asset_id(session, asset_id=asset.id)
|
||||
session.commit()
|
||||
|
||||
tags = get_asset_tags(session, asset_info_id=info.id)
|
||||
assert "missing" in tags
|
||||
|
||||
def test_add_missing_tag_is_idempotent(self, session: Session):
|
||||
asset = _make_asset(session, "hash1")
|
||||
info = _make_asset_info(session, asset)
|
||||
ensure_tags_exist(session, ["missing"], tag_type="system")
|
||||
|
||||
add_missing_tag_for_asset_id(session, asset_id=asset.id)
|
||||
add_missing_tag_for_asset_id(session, asset_id=asset.id)
|
||||
session.commit()
|
||||
|
||||
links = session.query(AssetInfoTag).filter_by(asset_info_id=info.id, tag_name="missing").all()
|
||||
assert len(links) == 1
|
||||
|
||||
def test_remove_missing_tag_for_asset_id(self, session: Session):
|
||||
asset = _make_asset(session, "hash1")
|
||||
info = _make_asset_info(session, asset)
|
||||
ensure_tags_exist(session, ["missing"], tag_type="system")
|
||||
add_missing_tag_for_asset_id(session, asset_id=asset.id)
|
||||
|
||||
remove_missing_tag_for_asset_id(session, asset_id=asset.id)
|
||||
session.commit()
|
||||
|
||||
tags = get_asset_tags(session, asset_info_id=info.id)
|
||||
assert "missing" not in tags
|
||||
|
||||
|
||||
class TestListTagsWithUsage:
|
||||
def test_returns_tags_with_counts(self, session: Session):
|
||||
ensure_tags_exist(session, ["used", "unused"])
|
||||
|
||||
asset = _make_asset(session, "hash1")
|
||||
info = _make_asset_info(session, asset)
|
||||
add_tags_to_asset_info(session, asset_info_id=info.id, tags=["used"])
|
||||
session.commit()
|
||||
|
||||
rows, total = list_tags_with_usage(session)
|
||||
|
||||
tag_dict = {name: count for name, _, count in rows}
|
||||
assert tag_dict["used"] == 1
|
||||
assert tag_dict["unused"] == 0
|
||||
assert total == 2
|
||||
|
||||
def test_exclude_zero_counts(self, session: Session):
|
||||
ensure_tags_exist(session, ["used", "unused"])
|
||||
|
||||
asset = _make_asset(session, "hash1")
|
||||
info = _make_asset_info(session, asset)
|
||||
add_tags_to_asset_info(session, asset_info_id=info.id, tags=["used"])
|
||||
session.commit()
|
||||
|
||||
rows, total = list_tags_with_usage(session, include_zero=False)
|
||||
|
||||
tag_names = {name for name, _, _ in rows}
|
||||
assert "used" in tag_names
|
||||
assert "unused" not in tag_names
|
||||
|
||||
def test_prefix_filter(self, session: Session):
|
||||
ensure_tags_exist(session, ["alpha", "beta", "alphabet"])
|
||||
session.commit()
|
||||
|
||||
rows, total = list_tags_with_usage(session, prefix="alph")
|
||||
|
||||
tag_names = {name for name, _, _ in rows}
|
||||
assert tag_names == {"alpha", "alphabet"}
|
||||
|
||||
def test_order_by_name(self, session: Session):
|
||||
ensure_tags_exist(session, ["zebra", "alpha", "middle"])
|
||||
session.commit()
|
||||
|
||||
rows, _ = list_tags_with_usage(session, order="name_asc")
|
||||
|
||||
names = [name for name, _, _ in rows]
|
||||
assert names == ["alpha", "middle", "zebra"]
|
||||
|
||||
def test_owner_visibility(self, session: Session):
|
||||
ensure_tags_exist(session, ["shared-tag", "owner-tag"])
|
||||
|
||||
asset = _make_asset(session, "hash1")
|
||||
shared_info = _make_asset_info(session, asset, name="shared", owner_id="")
|
||||
owner_info = _make_asset_info(session, asset, name="owned", owner_id="user1")
|
||||
|
||||
add_tags_to_asset_info(session, asset_info_id=shared_info.id, tags=["shared-tag"])
|
||||
add_tags_to_asset_info(session, asset_info_id=owner_info.id, tags=["owner-tag"])
|
||||
session.commit()
|
||||
|
||||
# Empty owner sees only shared
|
||||
rows, _ = list_tags_with_usage(session, owner_id="", include_zero=False)
|
||||
tag_dict = {name: count for name, _, count in rows}
|
||||
assert tag_dict.get("shared-tag", 0) == 1
|
||||
assert tag_dict.get("owner-tag", 0) == 0
|
||||
|
||||
# User1 sees both
|
||||
rows, _ = list_tags_with_usage(session, owner_id="user1", include_zero=False)
|
||||
tag_dict = {name: count for name, _, count in rows}
|
||||
assert tag_dict.get("shared-tag", 0) == 1
|
||||
assert tag_dict.get("owner-tag", 0) == 1
|
||||
Reference in New Issue
Block a user