import copy import unittest from unittest.mock import patch from step3_qa_agent.agent.plan_presentation import describe_plan from step3_qa_agent.agent import production SCHEMA = {'nodes':[{'id':'人员信息','attributes':['姓名','年龄','岗位']}, {'id':'项目信息','attributes':['名称']}], 'relations':[{'id':'r1','source':'人员信息','target':'项目信息','type':'服务于'}]} class PlanPresentationTests(unittest.TestCase): def plan(self): return {'steps':[{'step_id':'s1','tool':'图谱查询','depends':[], 'params':{ 'nodes':[{'alias':'n','type':'人员信息'}], 'filters':[{'alias':'n','field':'年龄','op':'lte','value':40}], 'select':[{'alias':'n','field':'姓名','as':'姓名'}], 'limit':20}}]} def test_conditions_and_confirmation_do_not_change_query(self): plan=self.plan(); before=copy.deepcopy(plan) state={'plan':plan,'qa_context':{'schema':SCHEMA,'build_id':'test'}} production.validate_plan(plan,state['qa_context']) with patch.object(production,'interrupt',return_value='确认') as interrupt: self.assertTrue(production.confirm(state)['plan_confirm']) text=interrupt.call_args.args[0]['message'] self.assertIn('年龄不大于“40”',text) self.assertIn('最多展示20条明细',text) self.assertNotIn('"nodes"',text) self.assertEqual(plan,before) with patch.object(production,'interrupt',return_value='改成30岁'): self.assertEqual(production.confirm(state)['plan_feedback'],'改成30岁') def test_relation_grouping_count_order(self): plan=self.plan(); p=plan['steps'][0]['params'] p['nodes'].append({'alias':'p','type':'项目信息'}) p['relations']=[{'source':'n','target':'p','id':'r1'}] p['select']=[{'alias':'p','field':'名称','as':'项目'}] p['aggregates']=[{'alias':'n','op':'count','as':'人数'}] p['order_by']=[{'field':'人数','direction':'desc'}] production.validate_plan(plan,{'schema':SCHEMA,'build_id':'test'}) text=describe_plan(plan,SCHEMA) for expected in ['服务于','按以下信息分组','同一对象只计一次','人数从高到低','20组统计结果']: self.assertIn(expected,text) def test_all_filter_operators_and_aggregate_forms(self): for op,value,expected in [('eq','甲','是“甲”'),('ne','甲','不是“甲”'),('contains','甲','包含“甲”'), ('in',['甲','乙'],'“甲”、“乙”'),('gt',0,'大于“0”'),('gte',0,'不小于“0”'), ('lt',40,'小于“40”'),('is_null',None,'未填写'),('not_null',None,'已填写')]: plan=self.plan(); plan['steps'][0]['params']['filters']=[{'alias':'n','field':'姓名','op':op,'value':value}] self.assertIn(expected,describe_plan(plan,SCHEMA)) for op,expected in [('sum','合计'),('avg','平均值'),('min','最小值'),('max','最大值'),('count_distinct','不同取值')]: plan=self.plan(); p=plan['steps'][0]['params']; p['select']=[]; p['filters']=[] p['aggregates']=[{'alias':'n','field':'年龄','op':op,'as':'统计'}] self.assertIn(expected,describe_plan(plan,SCHEMA)) self.assertIn('不额外限定',describe_plan(plan,SCHEMA)) def test_unsupported_does_not_interrupt(self): with patch.object(production,'interrupt') as interrupt: self.assertTrue(production.confirm({'plan':{'steps':[]}})['plan_confirm']) interrupt.assert_not_called() if __name__ == '__main__': unittest.main()