init
This commit is contained in:
@@ -0,0 +1,158 @@
|
||||
import json
|
||||
from unittest.mock import patch
|
||||
|
||||
from scripts.console_chat import format_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("mock answer")],
|
||||
"context": {
|
||||
"table_name": "apbo_eta_ful",
|
||||
"query_mode": "detail",
|
||||
"final_sql": "SELECT service_order_id FROM dwd_ai.apbo_eta_ful",
|
||||
"sql_plan": {"selected_table": "apbo_eta_ful", "query_mode": "detail"},
|
||||
},
|
||||
"final_step": "response_generated",
|
||||
}
|
||||
|
||||
|
||||
def test_format_result_includes_optional_blocks():
|
||||
result = {
|
||||
"messages": [FakeMessage("hello")],
|
||||
"context": {
|
||||
"final_sql": "SELECT 1",
|
||||
"sql_plan": {"mode": "detail"},
|
||||
"foo": "bar",
|
||||
},
|
||||
}
|
||||
|
||||
text = format_result(result, show_sql=True, show_context=True, show_plan=True)
|
||||
assert "Answer:" in text
|
||||
assert "SQL:" in text
|
||||
assert "SQL Plan:" in text
|
||||
assert "Context:" in text
|
||||
|
||||
|
||||
def test_format_result_renders_table_from_wrapped_sr_api_result():
|
||||
wrapped = 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,
|
||||
)
|
||||
result = {
|
||||
"messages": [FakeMessage('{"status_code": 200, "text": "..."}')],
|
||||
"context": {
|
||||
"sr_api_result": wrapped,
|
||||
"final_sql": "SELECT service_order_id, ship_to_country FROM dwd_ai.apbo_eta_ful",
|
||||
},
|
||||
}
|
||||
|
||||
text = format_result(result)
|
||||
assert "Query Result: 2 row(s)" in text
|
||||
assert "status=200" in text
|
||||
assert "service_order_id" in text
|
||||
assert "ship_to_country" in text
|
||||
assert "4020438779" in text
|
||||
assert "VN" in text
|
||||
|
||||
|
||||
def test_format_result_renders_columns_and_rows_payload():
|
||||
result = {
|
||||
"messages": [FakeMessage("ok")],
|
||||
"context": {
|
||||
"sr_api_result": {
|
||||
"status_code": 200,
|
||||
"text": {
|
||||
"columns": ["region", "qty"],
|
||||
"rows": [["ANZ", 12], ["CAP", 8]],
|
||||
},
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
text = format_result(result)
|
||||
assert "Query Result: 2 row(s)" in text
|
||||
assert "status=200" in text
|
||||
assert "region" in text
|
||||
assert "qty" in text
|
||||
assert "ANZ" in text
|
||||
assert "12" in text
|
||||
|
||||
|
||||
def test_format_result_falls_back_for_non_tabular_error():
|
||||
result = {
|
||||
"messages": [FakeMessage("请求失败: timeout")],
|
||||
"context": {"sr_api_result": "请求失败: timeout"},
|
||||
}
|
||||
|
||||
text = format_result(result)
|
||||
assert "Answer:" in text
|
||||
assert "请求失败: timeout" in text
|
||||
assert "Query Result:" not in text
|
||||
|
||||
|
||||
def test_main_one_shot_success(capsys):
|
||||
with patch("scripts.console_chat.Config.validate_config", return_value=None), \
|
||||
patch("scripts.console_chat.ConversationAgent", FakeAgent):
|
||||
exit_code = main([
|
||||
"--query",
|
||||
"查询 SO 4020438779 的 eta 信息",
|
||||
"--skip-sr-api",
|
||||
"--show-sql",
|
||||
])
|
||||
|
||||
captured = capsys.readouterr()
|
||||
assert exit_code == 0
|
||||
assert "mock answer" in captured.out
|
||||
assert "SELECT service_order_id FROM dwd_ai.apbo_eta_ful" in captured.out
|
||||
|
||||
|
||||
def test_run_turn_enables_debug_node_trace(capsys):
|
||||
agent = FakeAgent()
|
||||
|
||||
run_turn(
|
||||
agent,
|
||||
"查询 SO 4020438779 的 eta 信息",
|
||||
user="tester",
|
||||
conversation_id="cid-1",
|
||||
skip_sr_api=True,
|
||||
show_sql=False,
|
||||
show_context=False,
|
||||
show_plan=False,
|
||||
)
|
||||
|
||||
_, kwargs = agent.calls[-1]
|
||||
assert kwargs["debug_node_trace"] is True
|
||||
|
||||
|
||||
def test_main_config_error(capsys):
|
||||
with patch("scripts.console_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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user