""" 测试:对话内容提取器 ==================== 测试 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