98 lines
5.2 KiB
Python
98 lines
5.2 KiB
Python
# -*- coding: utf-8 -*-
|
|
import unittest
|
|
from unittest import mock
|
|
|
|
from flask import Flask, session
|
|
|
|
from routes import ai
|
|
from services.ai_datasources import AI_DATASOURCES
|
|
|
|
|
|
class AiAnalysisP0Test(unittest.TestCase):
|
|
def setUp(self):
|
|
self.app = Flask(__name__)
|
|
self.app.secret_key = 'test'
|
|
|
|
def test_normalize_tasks_accepts_only_registered_controlled_tools(self):
|
|
raw = [
|
|
{'datasource': 'product_profile', 'params': {'product_code': 'SKU-1', 'sql': 'DROP TABLE x', 'limit': 999}},
|
|
{'datasource': 'product_profile', 'params': {'brand': '重复'}},
|
|
{'datasource': 'not_exists', 'params': {}},
|
|
{'datasource': 'planning_metrics', 'params': {'limit': 'bad'}},
|
|
]
|
|
tasks = ai._normalize_tasks(raw)
|
|
self.assertEqual(['product_profile', 'planning_metrics'], [x['datasource'] for x in tasks])
|
|
self.assertNotIn('sql', tasks[0]['params'])
|
|
self.assertEqual(200, tasks[0]['params']['limit'])
|
|
self.assertEqual(50, tasks[1]['params']['limit'])
|
|
|
|
def test_chart_must_reference_real_dataset_fields(self):
|
|
datasets = [{'key': 'planning_metrics', 'columns': ['product_code', '建议下单数量']}]
|
|
valid = ai._validate_chart({
|
|
'type': 'bar', 'datasource': 'planning_metrics', 'label_field': 'product_code',
|
|
'value_fields': ['建议下单数量'], 'title': '建议下单排行'
|
|
}, datasets)
|
|
self.assertEqual('bar', valid['type'])
|
|
invalid = ai._validate_chart({
|
|
'type': 'bar', 'datasource': 'planning_metrics', 'label_field': '不存在',
|
|
'value_fields': ['建议下单数量']
|
|
}, datasets)
|
|
self.assertEqual('none', invalid['type'])
|
|
|
|
def test_multi_source_query_then_second_stage_analysis(self):
|
|
plan = {
|
|
'intent': 'analysis', 'title': 'SKU下单风险分析',
|
|
'tasks': [
|
|
{'datasource': 'product_profile', 'params': {'product_code': 'SKU-1'}, 'purpose': '基础资料'},
|
|
{'datasource': 'planning_metrics', 'params': {'product_code': 'SKU-1'}, 'purpose': '下单判断'},
|
|
]
|
|
}
|
|
analysis = {
|
|
'title': 'SKU下单风险分析', 'summary': '当前建议暂缓下单。',
|
|
'findings': [{'title': '库存', 'detail': '总支撑天数为40天', 'level': '信息'}],
|
|
'risks': [], 'actions': [{'priority': 'P0', 'action': '核对在途', 'reason': '避免重复下单'}],
|
|
'evidence': [{'label': '总支撑天数', 'value': '40天', 'source': '计划预测'}],
|
|
'confidence': '高', 'limitations': [], 'follow_ups': ['查看BOM包材是否够用'],
|
|
'chart': {'type': 'bar', 'datasource': 'planning_metrics', 'label_field': 'product_code',
|
|
'value_fields': ['建议下单数量'], 'title': '建议下单'}
|
|
}
|
|
fake_results = {
|
|
'product_profile': {'columns': ['product_code', '供应商'], 'rows': [{'product_code': 'SKU-1', '供应商': 'A'}]},
|
|
'planning_metrics': {'columns': ['product_code', '建议下单数量'], 'rows': [{'product_code': 'SKU-1', '建议下单数量': 0}]},
|
|
}
|
|
old_handlers = {key: AI_DATASOURCES[key]['handler'] for key in fake_results}
|
|
try:
|
|
for key, result in fake_results.items():
|
|
AI_DATASOURCES[key]['handler'] = mock.Mock(return_value=result)
|
|
with self.app.test_request_context('/api/ai/ask', method='POST', json={'question': 'SKU-1要不要下单'}):
|
|
session['user_id'] = 1
|
|
session['role'] = 'super_admin'
|
|
with mock.patch.object(ai, 'operation_required', return_value=True), \
|
|
mock.patch.object(ai, '_plan_analysis', return_value=plan), \
|
|
mock.patch.object(ai, '_synthesize_analysis', return_value=analysis), \
|
|
mock.patch.object(ai, '_can_view_datasource', return_value=True), \
|
|
mock.patch.object(ai, 'datasource_meta', side_effect=lambda key: {'key': key, 'name': AI_DATASOURCES[key]['name'], 'data_date': '2026-10-01'}), \
|
|
mock.patch.object(ai, '_save_history') as save_history, \
|
|
mock.patch.object(ai, 'add_operation_log'):
|
|
response = ai.ai_ask.__wrapped__()
|
|
payload = response.get_json()['data']
|
|
self.assertEqual('analysis', payload['type'])
|
|
self.assertEqual(2, len(payload['datasets']))
|
|
self.assertEqual('当前建议暂缓下单。', payload['summary'])
|
|
self.assertEqual('bar', payload['chart']['type'])
|
|
self.assertTrue(save_history.called)
|
|
self.assertEqual('analysis', save_history.call_args.args[4])
|
|
finally:
|
|
for key, handler in old_handlers.items():
|
|
AI_DATASOURCES[key]['handler'] = handler
|
|
|
|
def test_p0_datasources_cover_product_bom_order_and_supplier(self):
|
|
for key in ('product_profile', 'bom_supply_chain', 'order_detail', 'supplier_delivery'):
|
|
self.assertIn(key, AI_DATASOURCES)
|
|
self.assertTrue(callable(AI_DATASOURCES[key]['handler']))
|
|
self.assertTrue(AI_DATASOURCES[key].get('permission_pages'))
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|