life-echo/api/tests/test_classification_fragment.py

"""ClassificationAgent：零散档案启发式与 none→summary 兜底（纯函数/无 LLM）。"""

import pytest

from app.agents.memoir.classification_agent import (
    ClassificationAgent,
    _looks_like_fragment_only,
    _parse_category_from_llm_response,
)


@pytest.mark.parametrize(
    "text,expected_fragment",
    [
        ("", True),
        ("   ", True),
        ("我1999年出生", True),
        ("1999年出生。", True),
        ("1999年出生！", True),
        ("我是云南人", True),
        ("我是北京籍。", True),
        ("小学二年级那次下雨爷爷背我过河，鞋全湿了。", False),
        ("我出生在农村，家里养过一头黄牛。", False),
        ("我是北京人，后来去上海读了大学。", False),
    ],
)
def test_looks_like_fragment_only(text: str, expected_fragment: bool) -> None:
    assert _looks_like_fragment_only(text) is expected_fragment


def test_classify_maps_birth_year_fragment_to_summary_without_llm() -> None:
    agent = ClassificationAgent()
    result = agent.classify("1999年出生", fallback_stage="childhood", llm=None)
    assert result.category == "summary"
    assert result.llm_said_none is False


@pytest.mark.parametrize(
    "raw,expected",
    [
        ('{"category": "childhood"}', "childhood"),
        ('```json\n{"category": "none"}\n```', "none"),
        ("childhood", "childhood"),
        ('"education"', "education"),
    ],
)
def test_parse_category_from_llm_response(raw: str, expected: str) -> None:
    assert _parse_category_from_llm_response(raw) == expected


def test_classify_fallback_when_no_llm_and_narrative_snippet() -> None:
    agent = ClassificationAgent()
    out = agent.classify(
        "小学二年级的时候我在操场上摔了一跤，膝盖流了很多血，是老师背我去医务室的。",
        fallback_stage="childhood",
        llm=None,
    )
    assert out.category == "education"
    assert out.llm_said_none is False