kg-agent/scripts/test_plan_presentation.py

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()