knowledge-adapter-service/tests/test_weknora_resource_resolver.py

from contextlib import AbstractAsyncContextManager
from types import TracebackType
from uuid import UUID

import pytest

from knowledge_adapter.weknora_resource_resolver import (
    ResolvedAsset,
    ResourceResolutionError,
    WeKnoraResourceResolver,
)

HANDLE = "AbCdEfGhIjKlMnOpQrSt_1"
BACKEND_ID = "550e8400-e29b-41d4-a716-446655440000"


class FakeTransaction(AbstractAsyncContextManager[None]):
    def __init__(self) -> None:
        self.entered = False
        self.readonly: bool | None = None

    async def __aenter__(self) -> None:
        self.entered = True

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


class FakeConnection:
    def __init__(self, rows: list[dict[str, object]]) -> None:
        self.rows = rows
        self.transaction_context = FakeTransaction()
        self.query = ""
        self.args: tuple[object, ...] = ()

    def transaction(self, *, readonly: bool) -> FakeTransaction:
        self.transaction_context.readonly = readonly
        return self.transaction_context

    async def fetch(self, query: str, *args: object) -> list[dict[str, object]]:
        self.query = query
        self.args = args
        return self.rows


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

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

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


class FakePool:
    def __init__(self, connection: FakeConnection) -> None:
        self.connection = connection
        self.acquire_count = 0

    def acquire(self) -> AcquireContext:
        self.acquire_count += 1
        return AcquireContext(self.connection)


class FailingConnection(FakeConnection):
    async def fetch(self, query: str, *args: object) -> list[dict[str, object]]:
        raise RuntimeError(
            "postgresql://readonly:secret@db.internal resources physical_path=private/object"
        )


async def test_resolver_uses_minimal_parameterized_readonly_query() -> None:
    connection = FakeConnection(
        [
            {
                "tenant_id": 10008,
                "storage_backend_id": "backend-test",
                "provider": "oss",
                "physical_path": (
                    "storage://backend-test/"
                    "oss://private-assets/weknora/10008/assets/image.png"
                ),
                "mime_type": "image/png",
            }
        ]
    )
    resolver = WeKnoraResourceResolver(
        pool=FakePool(connection),
        tenant_id=10008,
        storage_backend_id="backend-test",
        bucket="private-assets",
    )

    resolved = await resolver.resolve(f"resource://{HANDLE}")

    assert resolved == ResolvedAsset(
        object_key="weknora/10008/assets/image.png",
        content_type="image/png",
    )
    assert connection.transaction_context.entered is True
    assert connection.transaction_context.readonly is True
    assert connection.args == (HANDLE, 10008, "backend-test")
    normalized_query = " ".join(connection.query.split()).lower()
    assert normalized_query.startswith(
        "select tenant_id, storage_backend_id, provider, physical_path, mime_type from resources"
    )
    assert "handle = $1" in normalized_query
    assert "tenant_id = $2" in normalized_query
    assert "storage_backend_id = $3" in normalized_query
    assert "state = 'active'" in normalized_query
    assert "deleted_at is null" in normalized_query
    for forbidden_column in ("content", "credentials", "api_key", "secret", "original_name"):
        assert forbidden_column not in normalized_query


async def test_resolver_accepts_asyncpg_uuid_storage_backend_id() -> None:
    connection = FakeConnection(
        [
            {
                "tenant_id": 10008,
                "storage_backend_id": UUID(BACKEND_ID),
                "provider": "oss",
                "physical_path": (
                    f"storage://{BACKEND_ID}/"
                    "oss://private-assets/weknora/10008/assets/image.png"
                ),
                "mime_type": "image/png",
            }
        ]
    )
    resolver = WeKnoraResourceResolver(
        pool=FakePool(connection),
        tenant_id=10008,
        storage_backend_id=BACKEND_ID,
        bucket="private-assets",
    )

    resolved = await resolver.resolve(f"resource://{HANDLE}")

    assert resolved.object_key == "weknora/10008/assets/image.png"


@pytest.mark.parametrize("row_count", [0, 2])
async def test_resolver_requires_exactly_one_resource_row(row_count: int) -> None:
    row = {
        "tenant_id": 10008,
        "storage_backend_id": "backend-test",
        "provider": "oss",
        "physical_path": (
            "storage://backend-test/"
            "oss://private-assets/weknora/10008/assets/image.png"
        ),
        "mime_type": "image/png",
    }
    resolver = WeKnoraResourceResolver(
        pool=FakePool(FakeConnection([row] * row_count)),
        tenant_id=10008,
        storage_backend_id="backend-test",
        bucket="private-assets",
    )

    with pytest.raises(ResourceResolutionError, match=r"^resource unavailable$"):
        await resolver.resolve(f"resource://{HANDLE}")


@pytest.mark.parametrize(
    "overrides",
    [
        {"tenant_id": 99999},
        {"storage_backend_id": "other-backend"},
        {"provider": "cos"},
    ],
)
async def test_resolver_rechecks_row_ownership(overrides: dict[str, object]) -> None:
    row = {
        "tenant_id": 10008,
        "storage_backend_id": "backend-test",
        "provider": "oss",
        "physical_path": (
            "storage://backend-test/"
            "oss://private-assets/weknora/10008/assets/image.png"
        ),
        "mime_type": "image/png",
        **overrides,
    }
    resolver = WeKnoraResourceResolver(
        pool=FakePool(FakeConnection([row])),
        tenant_id=10008,
        storage_backend_id="backend-test",
        bucket="private-assets",
    )

    with pytest.raises(ResourceResolutionError, match=r"^resource unavailable$"):
        await resolver.resolve(f"resource://{HANDLE}")


async def test_resolver_maps_database_error_to_safe_fixed_error() -> None:
    resolver = WeKnoraResourceResolver(
        pool=FakePool(FailingConnection([])),
        tenant_id=10008,
        storage_backend_id="backend-test",
        bucket="private-assets",
    )

    with pytest.raises(ResourceResolutionError) as captured:
        await resolver.resolve(f"resource://{HANDLE}")

    assert str(captured.value) == "resource unavailable"
    assert "postgresql" not in str(captured.value)
    assert "physical_path" not in str(captured.value)


async def test_invalid_resource_reference_is_rejected_before_database_access() -> None:
    pool = FakePool(FakeConnection([]))
    resolver = WeKnoraResourceResolver(
        pool=pool,
        tenant_id=10008,
        storage_backend_id="backend-test",
        bucket="private-assets",
    )

    with pytest.raises(ResourceResolutionError, match=r"^resource unavailable$"):
        await resolver.resolve("https://attacker.example/image.png")

    assert pool.acquire_count == 0