167 lines
5.0 KiB
Python
167 lines
5.0 KiB
Python
import json
|
|
from unittest.mock import patch
|
|
|
|
from scripts.demo_chat import format_demo_result, main, run_turn
|
|
|
|
|
|
class FakeMessage:
|
|
def __init__(self, content: str):
|
|
self.content = content
|
|
|
|
|
|
class FakeAgent:
|
|
def __init__(self, model_section=None):
|
|
self.model_section = model_section
|
|
self.calls = []
|
|
|
|
def run(self, query, **kwargs):
|
|
self.calls.append((query, kwargs))
|
|
return {
|
|
"messages": [FakeMessage('{"status_code": 200, "text": "..."}')],
|
|
"context": {
|
|
"final_sql": "SELECT service_order_id, ship_to_country FROM dwd_ai.apbo_eta_ful",
|
|
"sr_api_result": json.dumps(
|
|
{
|
|
"status_code": 200,
|
|
"text": json.dumps(
|
|
{
|
|
"data": [
|
|
{"service_order_id": "4020438779", "ship_to_country": "VN"},
|
|
{"service_order_id": "4020438780", "ship_to_country": "PH"},
|
|
]
|
|
},
|
|
ensure_ascii=False,
|
|
),
|
|
},
|
|
ensure_ascii=False,
|
|
),
|
|
},
|
|
"final_step": "response_generated",
|
|
}
|
|
|
|
|
|
def test_format_demo_result_includes_time_sql_row_count_and_table():
|
|
result = FakeAgent().run("查询")
|
|
|
|
text = format_demo_result(result, 1.234)
|
|
|
|
assert "耗时: 1.23s" in text
|
|
assert "SQL:" in text
|
|
assert "SELECT service_order_id, ship_to_country FROM dwd_ai.apbo_eta_ful" in text
|
|
assert "数据行数: 2" in text
|
|
assert "SQL执行结果表:" in text
|
|
assert "service_order_id" in text
|
|
assert "4020438779" in text
|
|
|
|
|
|
def test_format_demo_result_falls_back_when_result_is_not_tabular():
|
|
result = {
|
|
"messages": [FakeMessage("请求失败: timeout")],
|
|
"context": {
|
|
"final_sql": "SELECT 1",
|
|
"sr_api_result": "请求失败: timeout",
|
|
},
|
|
}
|
|
|
|
text = format_demo_result(result, 0.4)
|
|
|
|
assert "耗时: 0.40s" in text
|
|
assert "数据行数: 0" in text
|
|
assert "SQL执行结果:" in text
|
|
assert "请求失败: timeout" in text
|
|
|
|
|
|
def test_format_demo_result_renders_empty_structured_result_as_table():
|
|
result = {
|
|
"messages": [FakeMessage('{"status_code": 200, "text": "..."}')],
|
|
"context": {
|
|
"final_sql": "SELECT 1",
|
|
"sr_api_result": json.dumps(
|
|
{
|
|
"status_code": 200,
|
|
"text": json.dumps(
|
|
{"code": "0", "data": [], "msg": "操作成功", "total": 0},
|
|
ensure_ascii=False,
|
|
),
|
|
},
|
|
ensure_ascii=False,
|
|
),
|
|
},
|
|
}
|
|
|
|
text = format_demo_result(result, 0.4)
|
|
|
|
assert "数据行数: 0" in text
|
|
assert "SQL执行结果表:" in text
|
|
assert "<empty table>" in text
|
|
assert '"status_code": 200' not in text
|
|
|
|
|
|
def test_format_demo_result_prefers_fallback_answer_for_empty_result():
|
|
result = {
|
|
"messages": [FakeMessage("未查询到符合条件的数据，请尝试调整筛选条件。")],
|
|
"context": {
|
|
"final_sql": "SELECT 1",
|
|
"sr_api_result": json.dumps(
|
|
{
|
|
"status_code": 200,
|
|
"text": json.dumps(
|
|
{"code": "0", "data": [], "msg": "操作成功", "total": 0},
|
|
ensure_ascii=False,
|
|
),
|
|
},
|
|
ensure_ascii=False,
|
|
),
|
|
"is_empty_result": True,
|
|
"response_source": "model_empty_result_fallback",
|
|
},
|
|
}
|
|
|
|
text = format_demo_result(result, 0.4)
|
|
|
|
assert "结果说明:" in text
|
|
assert "未查询到符合条件的数据" in text
|
|
assert "SQL执行结果表:" not in text
|
|
|
|
|
|
def test_main_one_shot_success(capsys):
|
|
with patch("scripts.demo_chat.Config.validate_config", return_value=None), \
|
|
patch("scripts.demo_chat.ConversationAgent", FakeAgent):
|
|
exit_code = main([
|
|
"--query",
|
|
"查询 SO 4020438779 的 eta 信息",
|
|
])
|
|
|
|
captured = capsys.readouterr()
|
|
assert exit_code == 0
|
|
assert "耗时:" in captured.out
|
|
assert "数据行数: 2" in captured.out
|
|
assert "SELECT service_order_id, ship_to_country FROM dwd_ai.apbo_eta_ful" in captured.out
|
|
|
|
|
|
def test_run_turn_disables_debug_node_trace(capsys):
|
|
agent = FakeAgent()
|
|
|
|
run_turn(
|
|
agent,
|
|
"查询 SO 4020438779 的 eta 信息",
|
|
user="tester",
|
|
conversation_id="cid-1",
|
|
)
|
|
|
|
_, kwargs = agent.calls[-1]
|
|
assert kwargs["debug_node_trace"] is False
|
|
|
|
|
|
def test_main_config_error(capsys):
|
|
with patch("scripts.demo_chat.Config.validate_config", side_effect=ValueError("bad config")):
|
|
exit_code = main(["--query", "hello"])
|
|
|
|
captured = capsys.readouterr()
|
|
assert exit_code == 1
|
|
assert "Configuration error" in captured.err
|
|
|
|
|
|
|
|
|