knowledge-adapter-service/tests/test_asset_content_processor.py

import json
from datetime import UTC, datetime
from uuid import UUID

import pytest

from knowledge_adapter.asset_content_processor import (
    AssetContentProcessor,
    ContentProcessingError,
)
from knowledge_adapter.asset_repository import AssetMapping
from knowledge_adapter.models import WeKnoraChunk
from knowledge_adapter.weknora_resource_resolver import (
    ResolvedAsset,
    ResourceResolutionError,
)

SOURCE_REF = "resource://AbCdEfGhIjKlMnOpQrSt_1"
ASSET_ID = UUID("c7f61f89-19ce-45b1-ac60-dbeed5e2a491")
KB_ID = UUID("3904b40e-278b-4273-a6fa-db6cbc975587")


class RecordingAssetRepository:
    def __init__(
        self,
        mapping: AssetMapping | None,
        *,
        created_mapping: AssetMapping | None = None,
    ) -> None:
        self.mapping = mapping
        self.created_mapping = created_mapping
        self.lookup_calls: list[str] = []
        self.create_calls: list[tuple[str, str, str | None]] = []

    async def get_active_by_source_ref(self, source_ref: str) -> AssetMapping | None:
        self.lookup_calls.append(source_ref)
        return self.mapping

    async def get_or_create(
        self,
        *,
        source_ref: str,
        object_key: str,
        content_type: str | None,
    ) -> AssetMapping:
        self.create_calls.append((source_ref, object_key, content_type))
        if self.created_mapping is None:
            raise AssertionError("unexpected mapping creation")
        return self.created_mapping


class RecordingResolver:
    def __init__(
        self,
        resolved: ResolvedAsset | None = None,
        *,
        error: Exception | None = None,
    ) -> None:
        self.resolved = resolved
        self.error = error
        self.calls: list[str] = []

    async def resolve(self, source_ref: str) -> ResolvedAsset:
        self.calls.append(source_ref)
        if self.error is not None:
            raise self.error
        if self.resolved is None:
            raise AssertionError("unexpected WeKnora DB resolution")
        return self.resolved


def make_chunk(*, content: str, image_info: str = "") -> WeKnoraChunk:
    return WeKnoraChunk.model_validate(
        {
            "id": "09ab721f-7d05-45de-ae0f-dac303276f6d",
            "content": content,
            "score": 0.9,
            "knowledge_base_id": str(KB_ID),
            "knowledge_id": "6f77a258-27ff-4278-a667-8fe242ff8e10",
            "knowledge_title": "Manual",
            "knowledge_filename": "manual.pdf",
            "chunk_index": 0,
            "chunk_type": "text",
            "image_info": image_info,
        }
    )


def active_mapping() -> AssetMapping:
    return AssetMapping(
        asset_id=ASSET_ID,
        source_ref=SOURCE_REF,
        object_key="weknora/10008/assets/image.png",
        content_type="image/png",
        created_at=datetime(2026, 8, 24, tzinfo=UTC),
        status="active",
    )


async def test_markdown_resource_uses_existing_mapping_without_weknora_db() -> None:
    repository = RecordingAssetRepository(active_mapping())
    resolver = RecordingResolver()
    processor = AssetContentProcessor(
        repository=repository,
        resolver=resolver,
        public_base_url="https://assets.example.com",
    )
    chunk = make_chunk(content=f"Before ![diagram]({SOURCE_REF}) after")

    processed = await processor.process(chunk)

    assert processed == (
        f"Before ![diagram](https://assets.example.com/i/{ASSET_ID}) after"
    )
    assert repository.lookup_calls == [SOURCE_REF]
    assert repository.create_calls == []
    assert resolver.calls == []


async def test_html_img_resource_uses_existing_mapping_without_weknora_db() -> None:
    repository = RecordingAssetRepository(active_mapping())
    resolver = RecordingResolver()
    processor = AssetContentProcessor(
        repository=repository,
        resolver=resolver,
        public_base_url="https://assets.example.com",
    )
    chunk = make_chunk(content=f'<td><img src="{SOURCE_REF}"/></td>')

    processed = await processor.process(chunk)

    assert processed == (f'<td><img src="https://assets.example.com/i/{ASSET_ID}"/></td>')
    assert repository.lookup_calls == [SOURCE_REF]
    assert repository.create_calls == []
    assert resolver.calls == []


async def test_html_img_resource_preserves_attributes_and_quote_style() -> None:
    repository = RecordingAssetRepository(active_mapping())
    resolver = RecordingResolver()
    processor = AssetContentProcessor(
        repository=repository,
        resolver=resolver,
        public_base_url="https://assets.example.com",
    )
    chunk = make_chunk(content=f"<img alt='diagram' src='{SOURCE_REF}' width='100'>")

    processed = await processor.process(chunk)

    assert processed == (
        f"<img alt='diagram' src='https://assets.example.com/i/{ASSET_ID}' width='100'>"
    )
    assert repository.lookup_calls == [SOURCE_REF]
    assert repository.create_calls == []
    assert resolver.calls == []


async def test_unmapped_resource_is_resolved_and_persisted_once() -> None:
    repository = RecordingAssetRepository(
        None,
        created_mapping=active_mapping(),
    )
    resolver = RecordingResolver(
        ResolvedAsset(
            object_key="weknora/10008/assets/image.png",
            content_type="image/png",
        )
    )
    processor = AssetContentProcessor(
        repository=repository,
        resolver=resolver,
        public_base_url="https://assets.example.com",
    )

    processed = await processor.process(
        make_chunk(content=f"![diagram]({SOURCE_REF})")
    )

    assert processed == f"![diagram](https://assets.example.com/i/{ASSET_ID})"
    assert repository.lookup_calls == [SOURCE_REF]
    assert resolver.calls == [SOURCE_REF]
    assert repository.create_calls == [
        (SOURCE_REF, "weknora/10008/assets/image.png", "image/png")
    ]


async def test_duplicate_markdown_resource_is_deduplicated() -> None:
    repository = RecordingAssetRepository(active_mapping())
    resolver = RecordingResolver()
    processor = AssetContentProcessor(
        repository=repository,
        resolver=resolver,
        public_base_url="https://assets.example.com",
    )
    chunk = make_chunk(
        content=f"![first]({SOURCE_REF}) and ![second]({SOURCE_REF})"
    )

    processed = await processor.process(chunk)

    public_url = f"https://assets.example.com/i/{ASSET_ID}"
    assert processed == f"![first]({public_url}) and ![second]({public_url})"
    assert repository.lookup_calls == [SOURCE_REF]
    assert resolver.calls == []


async def test_image_info_only_resource_is_appended_as_stable_image() -> None:
    repository = RecordingAssetRepository(active_mapping())
    resolver = RecordingResolver()
    processor = AssetContentProcessor(
        repository=repository,
        resolver=resolver,
        public_base_url="https://assets.example.com",
    )
    image_info = json.dumps(
        [{"url": SOURCE_REF, "original_url": SOURCE_REF}],
        ensure_ascii=False,
    )

    processed = await processor.process(
        make_chunk(content="Text only", image_info=image_info)
    )

    public_url = f"https://assets.example.com/i/{ASSET_ID}"
    assert processed == f"Text only\n\n![知识图片]({public_url})"
    assert repository.lookup_calls == [SOURCE_REF]
    assert resolver.calls == []


async def test_malformed_image_info_fails_with_safe_error() -> None:
    processor = AssetContentProcessor(
        repository=RecordingAssetRepository(active_mapping()),
        resolver=RecordingResolver(),
        public_base_url="https://assets.example.com",
    )

    with pytest.raises(ContentProcessingError) as captured:
        await processor.process(make_chunk(content="Text", image_info="not-json"))

    assert str(captured.value) == "content processing failed"
    assert "not-json" not in str(captured.value)


async def test_non_array_image_info_fails_closed() -> None:
    processor = AssetContentProcessor(
        repository=RecordingAssetRepository(active_mapping()),
        resolver=RecordingResolver(),
        public_base_url="https://assets.example.com",
    )

    with pytest.raises(ContentProcessingError, match=r"^content processing failed$"):
        await processor.process(make_chunk(content="Text", image_info="{}"))


async def test_resource_resolution_error_is_mapped_to_content_error() -> None:
    processor = AssetContentProcessor(
        repository=RecordingAssetRepository(None),
        resolver=RecordingResolver(error=ResourceResolutionError("resource unavailable")),
        public_base_url="https://assets.example.com",
    )

    with pytest.raises(ContentProcessingError, match=r"^content processing failed$"):
        await processor.process(make_chunk(content=f"![diagram]({SOURCE_REF})"))


async def test_unrewritten_resource_reference_never_reaches_dify_content() -> None:
    processor = AssetContentProcessor(
        repository=RecordingAssetRepository(active_mapping()),
        resolver=RecordingResolver(),
        public_base_url="https://assets.example.com",
    )

    with pytest.raises(ContentProcessingError, match=r"^content processing failed$"):
        await processor.process(make_chunk(content=f"internal ref: {SOURCE_REF}"))


async def test_truncated_trailing_markdown_resource_image_is_removed() -> None:
    processor = AssetContentProcessor(
        repository=RecordingAssetRepository(active_mapping()),
        resolver=RecordingResolver(),
        public_base_url="https://assets.example.com",
    )

    processed = await processor.process(
        make_chunk(content="LED indicator meanings\n\n![diagram](resource://4j9P")
    )

    assert processed == "LED indicator meanings\n\n"


@pytest.mark.parametrize(
    "forbidden_content",
    [
        "local://srv/private/image.png",
        "/home/app/private/image.png",
        "http://127.0.0.1/image.png",
        "http://10.0.0.8/image.png",
        "http://192.168.1.8/image.png",
        "http://172.16.0.8/image.png",
        "https://private-assets.oss-cn-shenzhen-internal.aliyuncs.com/image.png",
        "https://private-assets.oss-cn-shenzhen.aliyuncs.com/image.png",
    ],
)
async def test_internal_locator_never_reaches_dify_content(
    forbidden_content: str,
) -> None:
    processor = AssetContentProcessor(
        repository=RecordingAssetRepository(active_mapping()),
        resolver=RecordingResolver(),
        public_base_url="https://assets.example.com",
    )

    with pytest.raises(ContentProcessingError, match=r"^content processing failed$"):
        await processor.process(make_chunk(content=forbidden_content))