tests/test_diet_extractor.py

"""
测试：对话内容提取器
====================
测试 DietExtractor 的正则提取和 fallback 机制
"""

import sys
sys.path.insert(0, '.')

from modules.diet_extractor import DietExtractor, get_extract_prompt


class TestDietExtractor:
    """测试饮食提取器"""

    def test_extract_diet_basic(self):
        """基本饮食记录提取"""
        text = """
[EXTRACT_DIET]
日期: 2026-06-08
早餐: 全麦面包, 牛奶
午餐: 番茄牛腩, 米饭
记录人: 丈夫
备注: 孕妇说想吃
[/EXTRACT_DIET]
"""
        result = DietExtractor.extract_all(text)
        assert len(result["diets"]) == 1
        diet = result["diets"][0]
        assert diet["日期"] == "2026-06-08"
        assert "午餐" in diet["meals"]
        assert "番茄牛腩" in diet["meals"]["午餐"]

    def test_extract_diet_multiple(self):
        """多条饮食记录提取"""
        text = """
[EXTRACT_DIET]
日期: 2026-06-08
早餐: 全麦面包, 牛奶
[/EXTRACT_DIET]
[EXTRACT_DIET]
日期: 2026-06-08
午餐: 清蒸鱼
[/EXTRACT_DIET]
"""
        result = DietExtractor.extract_all(text)
        assert len(result["diets"]) == 2

    def test_extract_diet_empty(self):
        """无饮食记录应返回空列表"""
        text = "今天天气真好"
        result = DietExtractor.extract_all(text)
        assert result["diets"] == []

    def test_extract_recipe(self):
        """菜谱提取"""
        text = """
[EXTRACT_RECIPE]
菜名: 清蒸鲈鱼
制作人: 丈夫
难度: 中等
食材: 鲈鱼, 葱, 姜
备注: 孕妇爱吃
[/EXTRACT_RECIPE]
"""
        result = DietExtractor.extract_all(text)
        assert len(result["recipes"]) == 1
        recipe = result["recipes"][0]
        assert recipe["菜名"] == "清蒸鲈鱼"
        assert recipe["制作人"] == "丈夫"

    def test_extract_preference(self):
        """偏好提取"""
        text = """
[EXTRACT_PREFERENCE]
人员: 小明
类型: 忌口
内容: 不吃香菜
[/EXTRACT_PREFERENCE]
"""
        result = DietExtractor.extract_all(text)
        assert len(result["preferences"]) == 1
        pref = result["preferences"][0]
        assert pref["人员"] == "小明"
        assert pref["类型"] == "忌口"

    def test_extract_weight(self):
        """体重提取"""
        text = """
[EXTRACT_WEIGHT]
日期: 2026-06-08
体重: 62.5
记录人: 孕妇
[/EXTRACT_WEIGHT]
"""
        result = DietExtractor.extract_all(text)
        assert len(result["weights"]) == 1
        w = result["weights"][0]
        assert w["体重"] == "62.5"

    def test_extract_memory(self):
        """家庭记忆提取"""
        text = """
[EXTRACT_MEMORY]
类型: 事件
内容: 丈夫今天做了番茄牛腩
[/EXTRACT_MEMORY]
"""
        result = DietExtractor.extract_all(text)
        assert len(result["memories"]) == 1
        m = result["memories"][0]
        assert m["类型"] == "事件"

    def test_extract_thinking(self):
        """思考步骤提取"""
        text = """
[THINKING]
当前步骤: 分析昨日钙摄入
下一步: 对比DRIs标准
[/THINKING]
"""
        result = DietExtractor.extract_all(text)
        assert result["thinking"] is not None
        assert result["thinking"]["当前步骤"] == "分析昨日钙摄入"
        assert result["thinking"]["下一步"] == "对比DRIs标准"

    def test_extract_all_types(self):
        """综合提取所有类型"""
        text = """
[EXTRACT_DIET]
日期: 2026-06-08
午餐: 番茄牛腩
[/EXTRACT_DIET]
[EXTRACT_RECIPE]
菜名: 番茄牛腩
制作人: 丈夫
[/EXTRACT_RECIPE]
[THINKING]
当前步骤: 分析
下一步: 推荐
[/THINKING]
"""
        result = DietExtractor.extract_all(text)
        assert len(result["diets"]) == 1
        assert len(result["recipes"]) == 1
        assert result["thinking"] is not None

    def test_fallback_extract_diet(self):
        """fallback 关键词匹配"""
        text = "我今天午餐吃了番茄牛腩和米饭"
        result = DietExtractor.fallback_extract_diet(text)
        assert result is not None
        assert "午餐" in result["meals"]

    def test_fallback_extract_thinking(self):
        """fallback 思考步骤匹配"""
        text = "当前步骤是分析钙摄入，下一步是推荐食物"
        result = DietExtractor.fallback_extract_thinking(text)
        assert result is not None

    def test_robust_extract_with_fallback(self):
        """综合提取：正则失败时用 fallback"""
        text = "我早餐吃了全麦面包和牛奶"
        result = DietExtractor.robust_extract(text)
        assert len(result["diets"]) > 0

    def test_robust_extract_with_markup(self):
        """综合提取：正则优先"""
        text = """
我吃了番茄牛腩
[EXTRACT_DIET]
日期: 2026-06-08
午餐: 番茄牛腩
[/EXTRACT_DIET]
"""
        result = DietExtractor.robust_extract(text)
        assert len(result["diets"]) == 1
        assert result["diets"][0].get("日期") == "2026-06-08"

    def test_get_extract_prompt(self):
        """提取提示词应包含日期"""
        prompt = get_extract_prompt("2026-06-08")
        assert "2026-06-08" in prompt
        assert "EXTRACT_DIET" in prompt
        assert "EXTRACT_RECIPE" in prompt
        assert "THINKING" in prompt