DataOps-platform/tests/knowledge/test_device_retrieval.py

from __future__ import annotations

from datetime import UTC, datetime

import pytest


def _row(**overrides):
    from app.core.knowledge.retrieval.device import DeviceSearchRow

    values = {
        "asset_uid": "11111111-1111-4111-8111-111111111111",
        "asset_type": "device",
        "name": "一号循环泵",
        "current_version": 3,
        "location": "动力车间",
        "organization": "设备动力部",
        "responsible_person": "张工",
        "source_codes": ("EQ-001", "ERP-PUMP-01"),
        "related_events": (
            ("fault", "轴承润滑故障", "FT-001"),
            ("alarm", "轴承温度持续升高", "AL-001"),
        ),
        "business_domain_uid": "22222222-2222-4222-8222-222222222222",
        "updated_at": datetime(2026, 7, 29, 8, 0, tzinfo=UTC),
        "rank": 0.95,
    }
    values.update(overrides)
    return DeviceSearchRow(**values)


def test_device_query_is_trimmed_bounded_and_required():
    from app.core.knowledge.retrieval.device import normalize_device_query

    assert normalize_device_query("  EQ-001  ") == "EQ-001"
    with pytest.raises(ValueError, match="不能为空"):
        normalize_device_query("   ")
    with pytest.raises(ValueError, match="300"):
        normalize_device_query("x" * 301)


@pytest.mark.parametrize(
    ("query", "expected_term"),
    [
        ("循环泵最近有哪些故障？", "循环泵"),
        ("动力车间有哪些设备？", "动力车间"),
        ("张工负责哪些设备？", "张工"),
        ("一号循环泵的责任人是谁？", "一号循环泵"),
        ("源 ID 为 EQ-001 的设备责任人是谁？", "EQ-001"),
        (
            "平台 UID：11111111-1111-4111-8111-111111111111",
            "11111111-1111-4111-8111-111111111111",
        ),
    ],
)
def test_device_query_terms_extract_searchable_entity_from_common_questions(
    query,
    expected_term,
):
    from app.core.knowledge.retrieval.device import device_query_terms

    terms = device_query_terms(query)

    assert terms[0] == expected_term
    assert normalize_device_query_for_assertion(query) in terms


def normalize_device_query_for_assertion(query):
    from app.core.knowledge.retrieval.device import normalize_device_query

    return normalize_device_query(query)


def test_device_evidence_is_stable_readable_and_does_not_expose_raw_payloads():
    from app.core.knowledge.retrieval.device import build_device_evidence

    evidence = build_device_evidence(_row())

    assert evidence.chunk_id == (
        "device:11111111-1111-4111-8111-111111111111:v3"
    )
    assert evidence.object_type == "DeviceAsset"
    assert evidence.object_version == 3
    assert evidence.point_keys == (
        "DeviceAsset/11111111-1111-4111-8111-111111111111/summary",
    )
    assert evidence.point_revisions == (3,)
    assert evidence.score == 0.95
    assert "设备名称：一号循环泵" in evidence.content
    assert "源 ID：EQ-001、ERP-PUMP-01" in evidence.content
    assert "位置：动力车间" in evidence.content
    assert "责任人：张工" in evidence.content
    assert "故障：轴承润滑故障（FT-001）" in evidence.content
    assert "告警：轴承温度持续升高（AL-001）" in evidence.content
    for forbidden in ("password", "permission_scope", "source_config"):
        assert forbidden not in evidence.content


def test_device_evidence_is_deterministic_for_unsorted_duplicate_inputs():
    from app.core.knowledge.retrieval.device import build_device_evidence

    first = build_device_evidence(
        _row(
            source_codes=("ERP-PUMP-01", "EQ-001", "EQ-001"),
            related_events=(
                ("alarm", "轴承温度持续升高", "AL-001"),
                ("fault", "轴承润滑故障", "FT-001"),
                ("fault", "轴承润滑故障", "FT-001"),
            ),
        )
    )
    second = build_device_evidence(_row())

    assert first.content == second.content
    assert first.point_keys == second.point_keys


def test_device_result_limit_is_bounded_before_repository_query():
    from app.core.knowledge.access import KnowledgeAccessContext
    from app.core.knowledge.retrieval.device import SqlDeviceKnowledgeRetriever

    class Repository:
        def __init__(self):
            self.calls = []

        def search(self, **kwargs):
            self.calls.append(kwargs)
            return [_row()]

    repository = Repository()
    retriever = SqlDeviceKnowledgeRetriever(repository)
    context = KnowledgeAccessContext(
        subject_id="viewer-1",
        roles=frozenset({"viewer"}),
        permissions=frozenset({"governance:read"}),
        business_domain_uids=frozenset(
            {"22222222-2222-4222-8222-222222222222"}
        ),
        correlation_id="correlation-1",
    )

    result = retriever.retrieve("循环泵", context, 999)

    assert len(result) == 1
    assert repository.calls == [
        {
            "query": "循环泵",
            "global_access": False,
            "business_domain_uids": (
                "22222222-2222-4222-8222-222222222222",
            ),
            "limit": 100,
        }
    ]