dify/dify-agent/tests/local/dify_agent/layers/drive/test_layer.py

"""Behavior tests for the runtime Dify drive layer."""

from __future__ import annotations

from pathlib import Path
import pytest

from agenton.layers import EmptyRuntimeState, LayerConfig, NoLayerDeps, PlainLayer
from dify_agent.layers.drive import DifyDriveLayerConfig, DifyDriveSkillConfig
from dify_agent.layers.drive.layer import DifyDriveLayer, DifyDriveLayerError, _DriveManifestItem


class _FakeExecutionContextConfig(LayerConfig):
    tenant_id: str


class _FakeExecutionContextLayer(PlainLayer[NoLayerDeps, _FakeExecutionContextConfig, EmptyRuntimeState]):
    type_id = None

    def __init__(self, tenant_id: str) -> None:
        self.config = _FakeExecutionContextConfig(tenant_id=tenant_id)


def _build_layer(tmp_path: Path) -> DifyDriveLayer:
    layer = DifyDriveLayer.from_config_with_settings(
        DifyDriveLayerConfig(
            drive_ref="agent-1",
            skills=[
                DifyDriveSkillConfig(
                    path="tender-analyzer",
                    name="Tender Analyzer",
                    description="Parses RFPs.",
                    skill_md_key="tender-analyzer/SKILL.md",
                    archive_key="tender-analyzer/.DIFY-SKILL-FULL.zip",
                ),
                DifyDriveSkillConfig(
                    path="other-skill",
                    name="Other Skill",
                    description="Fallback catalog entry.",
                    skill_md_key="other-skill/SKILL.md",
                    archive_key=None,
                ),
            ],
            mentioned_skill_keys=["tender-analyzer/SKILL.md"],
            mentioned_file_keys=["files/report.pdf"],
        ),
        inner_api_url="https://api.example.com",
        inner_api_key="secret",
    )
    layer.bind_deps({"execution_context": _FakeExecutionContextLayer("tenant-1")})
    return layer


@pytest.mark.anyio
async def test_on_context_create_loads_mentioned_targets_into_prompt(
    monkeypatch: pytest.MonkeyPatch, tmp_path: Path
) -> None:
    layer = _build_layer(tmp_path)

    async def _fetch_manifest_items(*, tenant_id: str, targets: list[tuple[str, bool]]) -> list[_DriveManifestItem]:
        assert tenant_id == "tenant-1"
        assert targets == [("tender-analyzer/", False), ("files/report.pdf", True)]
        return [
            _DriveManifestItem(key="tender-analyzer/SKILL.md", download_url="https://files/skill-md"),
            _DriveManifestItem(key="files/report.pdf", download_url="https://files/report"),
        ]

    async def _download_items(items: list[_DriveManifestItem]) -> dict[str, str]:
        assert {item.key for item in items} == {"files/report.pdf", "tender-analyzer/SKILL.md"}
        skill_path = tmp_path / "tender-analyzer" / "SKILL.md"
        skill_path.parent.mkdir(parents=True, exist_ok=True)
        skill_path.write_text("# Tender Analyzer\nUse carefully.\n", encoding="utf-8")
        file_path = tmp_path / "files" / "report.pdf"
        file_path.parent.mkdir(parents=True, exist_ok=True)
        file_path.write_bytes(b"pdf")
        return {
            "tender-analyzer/SKILL.md": str(skill_path),
            "files/report.pdf": str(file_path),
        }

    monkeypatch.setattr(layer, "_fetch_manifest_items", _fetch_manifest_items)
    monkeypatch.setattr(layer, "_download_items", _download_items)

    await layer.on_context_create()

    prompt = layer.build_prompt_context()
    assert "Loaded mentioned skills" in prompt
    assert "# Tender Analyzer\nUse carefully." in prompt
    assert f"files/report.pdf -> {tmp_path / 'files' / 'report.pdf'}" in prompt
    assert "Other available skills" in prompt
    assert "other-skill: Other Skill — Fallback catalog entry." in prompt


@pytest.mark.anyio
async def test_on_context_resume_loads_mentioned_targets_into_prompt(
    monkeypatch: pytest.MonkeyPatch, tmp_path: Path
) -> None:
    layer = _build_layer(tmp_path)

    async def _fetch_manifest_items(*, tenant_id: str, targets: list[tuple[str, bool]]) -> list[_DriveManifestItem]:
        assert tenant_id == "tenant-1"
        assert targets == [("tender-analyzer/", False), ("files/report.pdf", True)]
        return [
            _DriveManifestItem(key="tender-analyzer/SKILL.md", download_url="https://files/skill-md"),
            _DriveManifestItem(key="files/report.pdf", download_url="https://files/report"),
        ]

    async def _download_items(items: list[_DriveManifestItem]) -> dict[str, str]:
        assert {item.key for item in items} == {"files/report.pdf", "tender-analyzer/SKILL.md"}
        skill_path = tmp_path / "tender-analyzer" / "SKILL.md"
        skill_path.parent.mkdir(parents=True, exist_ok=True)
        skill_path.write_text("# Tender Analyzer\nUse carefully.\n", encoding="utf-8")
        file_path = tmp_path / "files" / "report.pdf"
        file_path.parent.mkdir(parents=True, exist_ok=True)
        file_path.write_bytes(b"pdf")
        return {
            "tender-analyzer/SKILL.md": str(skill_path),
            "files/report.pdf": str(file_path),
        }

    monkeypatch.setattr(layer, "_fetch_manifest_items", _fetch_manifest_items)
    monkeypatch.setattr(layer, "_download_items", _download_items)

    await layer.on_context_resume()

    prompt = layer.build_prompt_context()
    assert "Loaded mentioned skills" in prompt
    assert "# Tender Analyzer\nUse carefully." in prompt
    assert f"files/report.pdf -> {tmp_path / 'files' / 'report.pdf'}" in prompt
    assert "Other available skills" in prompt
    assert "other-skill: Other Skill — Fallback catalog entry." in prompt


@pytest.mark.anyio
async def test_on_context_create_raises_when_mentioned_file_is_missing(
    monkeypatch: pytest.MonkeyPatch,
    tmp_path: Path,
) -> None:
    layer = _build_layer(tmp_path)

    async def _fetch_manifest_items(*, tenant_id: str, targets: list[tuple[str, bool]]) -> list[_DriveManifestItem]:
        del tenant_id, targets
        return [_DriveManifestItem(key="tender-analyzer/SKILL.md", download_url="https://files/skill-md")]

    async def _download_items(items: list[_DriveManifestItem]) -> dict[str, str]:
        del items
        skill_path = tmp_path / "tender-analyzer" / "SKILL.md"
        skill_path.parent.mkdir(parents=True, exist_ok=True)
        skill_path.write_text("# Tender Analyzer\nUse carefully.\n", encoding="utf-8")
        return {"tender-analyzer/SKILL.md": str(skill_path)}

    monkeypatch.setattr(layer, "_fetch_manifest_items", _fetch_manifest_items)
    monkeypatch.setattr(layer, "_download_items", _download_items)

    with pytest.raises(DifyDriveLayerError, match="missing pulled file"):
        await layer.on_context_create()


@pytest.mark.anyio
async def test_on_context_resume_raises_when_mentioned_targets_are_missing(
    monkeypatch: pytest.MonkeyPatch,
    tmp_path: Path,
) -> None:
    layer = _build_layer(tmp_path)

    async def _fetch_manifest_items(*, tenant_id: str, targets: list[tuple[str, bool]]) -> list[_DriveManifestItem]:
        del tenant_id, targets
        return []

    async def _download_items(items: list[_DriveManifestItem]) -> dict[str, str]:
        assert items == []
        return {}

    monkeypatch.setattr(layer, "_fetch_manifest_items", _fetch_manifest_items)
    monkeypatch.setattr(layer, "_download_items", _download_items)

    with pytest.raises(DifyDriveLayerError, match="missing pulled file"):
        await layer.on_context_resume()


@pytest.mark.anyio
async def test_on_context_create_raises_when_manifest_is_empty_for_mentioned_targets(
    monkeypatch: pytest.MonkeyPatch,
    tmp_path: Path,
) -> None:
    layer = _build_layer(tmp_path)

    async def _fetch_manifest_items(*, tenant_id: str, targets: list[tuple[str, bool]]) -> list[_DriveManifestItem]:
        del tenant_id, targets
        return []

    async def _download_items(items: list[_DriveManifestItem]) -> dict[str, str]:
        assert items == []
        return {}

    monkeypatch.setattr(layer, "_fetch_manifest_items", _fetch_manifest_items)
    monkeypatch.setattr(layer, "_download_items", _download_items)

    with pytest.raises(DifyDriveLayerError, match="missing pulled file"):
        await layer.on_context_create()