PlantAI-dev/tests/unit/test_intent_species_roles.py

import pytest

from gfa_core.mas.agents.intent_v2 import IntentAgentV2
from gfa_core.mas.utils.intent_parser import (
    parse_explicit_species_role_mentions,
    parse_gene_family_scope_selection,
)


def test_ref_data_path_does_not_assign_species_roles() -> None:
    query = (
        "苦荞基因组数据在 "
        "/home/xia-lab/wutx_old/wutx/plantAnalysisAgent/ref_data/Tartary_genome，"
        "拟南芥基因组数据在 "
        "/home/xia-lab/wutx_old/wutx/plantAnalysisAgent/ref_data/ath，"
        "Ft是目标物种，At是参考物种"
    )

    scope = parse_gene_family_scope_selection(query)

    assert scope["targetSpecies"] == ["Ft"]
    assert scope["referenceSpecies"] == ["At"]


def test_explicit_role_parser_keeps_unassigned_mentions_separate() -> None:
    roles = parse_explicit_species_role_mentions(
        "Os和At的数据已经准备好；Os是目标物种，At是参考物种，Ppy稍后再分析"
    )

    assert roles == {
        "targetSpecies": ["Os"],
        "referenceSpecies": ["At"],
        "comparativeSpecies": [],
        "unassigned": ["Ppy"],
    }


@pytest.mark.parametrize("species", ["Gm", "Br", "Pt", "Sb"])
def test_scope_parser_preserves_species_supported_by_legacy_intent(species: str) -> None:
    scope = parse_gene_family_scope_selection(f"分析 {species} 的 SPL 基因家族")

    assert scope["targetSpecies"] == [species]


def test_scope_parser_accepts_explicit_role_bound_custom_species_code() -> None:
    scope = parse_gene_family_scope_selection("目标物种 Sv；参考物种 At")

    assert scope["targetSpecies"] == ["Sv"]
    assert scope["referenceSpecies"] == ["At"]


def test_scope_parser_does_not_treat_gene_family_as_species() -> None:
    scope = parse_gene_family_scope_selection("目标物种 Ft；分析 SPL 基因家族")

    assert scope["targetSpecies"] == ["Ft"]


def test_scope_parser_does_not_bind_custom_code_to_generic_goal_word() -> None:
    scope = parse_gene_family_scope_selection("分析目标 SPL 基因家族")

    assert scope["targetSpecies"] == []


def test_intent_validation_keeps_confirmed_target_and_reference_separate() -> None:
    agent = object.__new__(IntentAgentV2)

    intent = agent._validate_intent(
        {"gene_family": "SPL", "stages": ["identify"], "params": {}},
        (
            "请按照你规划的步骤来帮我执行\n"
            "已确认上下文：基因家族 SPL；目标物种 Ft；参考物种 At。"
        ),
    )

    assert intent["target_species"] == ["Ft"]
    assert intent["reference_species"] == ["At"]
    assert intent["species"] == "Ft"


def test_intent_validation_preserves_unlabelled_multi_target_request() -> None:
    agent = object.__new__(IntentAgentV2)

    intent = agent._validate_intent(
        {"gene_family": "SPL", "stages": ["identify"], "params": {}},
        "分析 Ft 和 Os 的 SPL 基因家族",
    )

    assert intent["target_species"] == ["Ft", "Os"]
    assert intent["reference_species"] == []
    assert intent["species"] == "Multi"