knowledge-adapter-service/tests/test_asset_repository.py

from contextlib import AbstractAsyncContextManager
from datetime import UTC, datetime
from types import TracebackType
from uuid import UUID

import pytest

from knowledge_adapter.asset_repository import (
    AssetMapping,
    AssetMappingConflictError,
    AssetRepository,
    AssetRepositoryError,
)

SOURCE_REF = "resource://AbCdEfGhIjKlMnOpQrSt_1"
ASSET_ID = UUID("c7f61f89-19ce-45b1-ac60-dbeed5e2a491")
CREATED_AT = datetime(2026, 8, 24, tzinfo=UTC)


class FakeAssetConnection:
    def __init__(self, row: dict[str, object] | None) -> None:
        self.row = row
        self.query = ""
        self.args: tuple[object, ...] = ()

    async def fetchrow(self, query: str, *args: object) -> dict[str, object] | None:
        self.query = query
        self.args = args
        return self.row


class FailingAssetConnection(FakeAssetConnection):
    async def fetchrow(self, query: str, *args: object) -> dict[str, object] | None:
        raise RuntimeError("postgresql://adapter:secret@db object_key=private/image.png")


class AcquireAssetConnection(AbstractAsyncContextManager[FakeAssetConnection]):
    def __init__(self, connection: FakeAssetConnection) -> None:
        self.connection = connection

    async def __aenter__(self) -> FakeAssetConnection:
        return self.connection

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc_value: BaseException | None,
        traceback: TracebackType | None,
    ) -> None:
        return None


class FakeAssetPool:
    def __init__(self, connection: FakeAssetConnection) -> None:
        self.connection = connection

    def acquire(self) -> AcquireAssetConnection:
        return AcquireAssetConnection(self.connection)


async def test_get_or_create_uses_atomic_source_ref_upsert() -> None:
    connection = FakeAssetConnection(
        {
            "asset_id": ASSET_ID,
            "source_ref": SOURCE_REF,
            "object_key": "weknora/10008/assets/image.png",
            "content_type": "image/png",
            "created_at": CREATED_AT,
            "status": "active",
        }
    )
    repository = AssetRepository(pool=FakeAssetPool(connection))

    mapping = await repository.get_or_create(
        source_ref=SOURCE_REF,
        object_key="weknora/10008/assets/image.png",
        content_type="image/png",
    )

    assert mapping == AssetMapping(
        asset_id=ASSET_ID,
        source_ref=SOURCE_REF,
        object_key="weknora/10008/assets/image.png",
        content_type="image/png",
        created_at=CREATED_AT,
        status="active",
    )
    generated_asset_id = connection.args[0]
    assert isinstance(generated_asset_id, UUID)
    assert generated_asset_id.version == 4
    assert connection.args[1:] == (
        SOURCE_REF,
        "weknora/10008/assets/image.png",
        "image/png",
    )
    sql = " ".join(connection.query.lower().split())
    assert "insert into asset_mapping" in sql
    assert "(asset_id, source_ref, object_key, content_type, status)" in sql
    assert "values ($1, $2, $3, $4, 'active')" in sql
    assert "on conflict (source_ref) do update" in sql
    assert "source_ref = asset_mapping.source_ref" in sql
    assert "where asset_mapping.object_key = excluded.object_key" in sql
    assert (
        "returning asset_id, source_ref, object_key, content_type, created_at, status" in sql
    )


async def test_get_or_create_rejects_source_ref_object_key_conflict() -> None:
    repository = AssetRepository(pool=FakeAssetPool(FakeAssetConnection(None)))

    with pytest.raises(AssetMappingConflictError, match=r"^asset mapping conflict$"):
        await repository.get_or_create(
            source_ref=SOURCE_REF,
            object_key="different/object.png",
            content_type="image/png",
        )


async def test_get_or_create_maps_database_error_to_safe_error() -> None:
    repository = AssetRepository(
        pool=FakeAssetPool(FailingAssetConnection(None)),
    )

    with pytest.raises(AssetRepositoryError) as captured:
        await repository.get_or_create(
            source_ref=SOURCE_REF,
            object_key="weknora/10008/assets/image.png",
            content_type="image/png",
        )

    assert str(captured.value) == "asset repository unavailable"
    assert "postgresql" not in str(captured.value)
    assert "object_key" not in str(captured.value)


async def test_get_active_by_source_ref_reads_only_active_mapping() -> None:
    connection = FakeAssetConnection(
        {
            "asset_id": ASSET_ID,
            "source_ref": SOURCE_REF,
            "object_key": "weknora/10008/assets/image.png",
            "content_type": "image/png",
            "created_at": CREATED_AT,
            "status": "active",
        }
    )
    repository = AssetRepository(pool=FakeAssetPool(connection))

    mapping = await repository.get_active_by_source_ref(SOURCE_REF)

    assert mapping is not None
    assert mapping.asset_id == ASSET_ID
    assert connection.args == (SOURCE_REF,)
    sql = " ".join(connection.query.lower().split())
    assert sql.startswith(
        "select asset_id, source_ref, object_key, content_type, created_at, status "
        "from asset_mapping"
    )
    assert "source_ref = $1" in sql
    assert "status = 'active'" in sql


async def test_get_active_by_source_ref_returns_none_when_missing() -> None:
    repository = AssetRepository(pool=FakeAssetPool(FakeAssetConnection(None)))

    mapping = await repository.get_active_by_source_ref(SOURCE_REF)

    assert mapping is None


async def test_get_active_by_source_ref_maps_database_error_to_safe_error() -> None:
    repository = AssetRepository(
        pool=FakeAssetPool(FailingAssetConnection(None)),
    )

    with pytest.raises(AssetRepositoryError, match=r"^asset repository unavailable$"):
        await repository.get_active_by_source_ref(SOURCE_REF)


async def test_get_active_by_asset_id_reads_only_active_mapping() -> None:
    connection = FakeAssetConnection(
        {
            "asset_id": ASSET_ID,
            "source_ref": SOURCE_REF,
            "object_key": "weknora/10008/assets/image.png",
            "content_type": "image/png",
            "created_at": CREATED_AT,
            "status": "active",
        }
    )
    repository = AssetRepository(pool=FakeAssetPool(connection))

    mapping = await repository.get_active_by_asset_id(ASSET_ID)

    assert mapping is not None
    assert mapping.object_key == "weknora/10008/assets/image.png"
    assert connection.args == (ASSET_ID,)
    sql = " ".join(connection.query.lower().split())
    assert "from asset_mapping" in sql
    assert "asset_id = $1" in sql
    assert "status = 'active'" in sql


async def test_get_active_by_asset_id_returns_none_when_missing() -> None:
    repository = AssetRepository(pool=FakeAssetPool(FakeAssetConnection(None)))

    mapping = await repository.get_active_by_asset_id(ASSET_ID)

    assert mapping is None