234 lines
9.2 KiB
Python
234 lines
9.2 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
Unit tests for Phase 3: Agent Changes
|
||
|
||
Run: python -m pytest tests/test_phase3.py -v
|
||
"""
|
||
|
||
import os
|
||
import sys
|
||
import json
|
||
|
||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||
|
||
import pytest
|
||
from data_analysis_agent import DataAnalysisAgent
|
||
from prompts import data_analysis_system_prompt, final_report_system_prompt
|
||
|
||
|
||
# ===========================================================================
|
||
# Task 8.1: _summarize_result
|
||
# ===========================================================================
|
||
|
||
class TestSummarizeResult:
|
||
@pytest.fixture
|
||
def agent(self):
|
||
"""Create a minimal DataAnalysisAgent for testing."""
|
||
agent = DataAnalysisAgent.__new__(DataAnalysisAgent)
|
||
agent._session_ref = None
|
||
return agent
|
||
|
||
def test_success_with_evidence_rows(self, agent):
|
||
"""8.1: Success with evidence rows produces DataFrame summary."""
|
||
result = {
|
||
"success": True,
|
||
"evidence_rows": [{"a": 1, "b": 2}, {"a": 3, "b": 4}],
|
||
"auto_exported_files": [{"variable_name": "df", "filename": "df.csv", "rows": 150, "cols": 8, "columns": []}],
|
||
}
|
||
summary = agent._summarize_result(result)
|
||
assert "执行成功" in summary
|
||
assert "DataFrame" in summary
|
||
assert "150" in summary
|
||
assert "8" in summary
|
||
|
||
def test_success_with_evidence_no_auto_files(self, agent):
|
||
"""8.1: Success with evidence but no auto_exported_files uses evidence length."""
|
||
result = {
|
||
"success": True,
|
||
"evidence_rows": [{"x": 1}, {"x": 2}, {"x": 3}],
|
||
"auto_exported_files": [],
|
||
}
|
||
summary = agent._summarize_result(result)
|
||
assert "执行成功" in summary
|
||
assert "DataFrame" in summary
|
||
|
||
def test_success_with_output(self, agent):
|
||
"""8.1: Success with output but no evidence shows first line."""
|
||
result = {
|
||
"success": True,
|
||
"evidence_rows": [],
|
||
"output": "Hello World\nSecond line",
|
||
}
|
||
summary = agent._summarize_result(result)
|
||
assert "执行成功" in summary
|
||
assert "Hello World" in summary
|
||
|
||
def test_success_no_output(self, agent):
|
||
"""8.1: Success with no output or evidence."""
|
||
result = {"success": True, "evidence_rows": [], "output": ""}
|
||
summary = agent._summarize_result(result)
|
||
assert summary == "执行成功"
|
||
|
||
def test_failure_short_error(self, agent):
|
||
"""8.1: Failure with short error message."""
|
||
result = {"success": False, "error": "KeyError: 'col_x'"}
|
||
summary = agent._summarize_result(result)
|
||
assert "执行失败" in summary
|
||
assert "KeyError" in summary
|
||
|
||
def test_failure_long_error_truncated(self, agent):
|
||
"""8.1: Failure with long error is truncated to 100 chars."""
|
||
long_error = "A" * 200
|
||
result = {"success": False, "error": long_error}
|
||
summary = agent._summarize_result(result)
|
||
assert "执行失败" in summary
|
||
assert "..." in summary
|
||
# The error portion should be at most 103 chars (100 + "...")
|
||
error_part = summary.split("执行失败: ")[1]
|
||
assert len(error_part) <= 104
|
||
|
||
def test_failure_no_error_field(self, agent):
|
||
"""8.1: Failure with missing error field."""
|
||
result = {"success": False}
|
||
summary = agent._summarize_result(result)
|
||
assert "执行失败" in summary
|
||
|
||
|
||
# ===========================================================================
|
||
# Task 8.2-8.4: Round_Data construction and session integration
|
||
# ===========================================================================
|
||
|
||
class TestRoundDataConstruction:
|
||
def test_handle_generate_code_returns_reasoning(self):
|
||
"""8.2: _handle_generate_code returns reasoning from yaml_data."""
|
||
agent = DataAnalysisAgent.__new__(DataAnalysisAgent)
|
||
agent._session_ref = None
|
||
# We need a minimal executor mock
|
||
from unittest.mock import MagicMock
|
||
agent.executor = MagicMock()
|
||
agent.executor.execute_code.return_value = {
|
||
"success": True, "output": "ok", "error": "",
|
||
"variables": {}, "evidence_rows": [],
|
||
"auto_exported_files": [], "prompt_saved_files": [],
|
||
}
|
||
yaml_data = {"code": "x = 1", "reasoning": "Testing reasoning field"}
|
||
result = agent._handle_generate_code("response text", yaml_data)
|
||
assert result["reasoning"] == "Testing reasoning field"
|
||
|
||
def test_handle_generate_code_empty_reasoning(self):
|
||
"""8.2: _handle_generate_code returns empty reasoning when not in yaml_data."""
|
||
agent = DataAnalysisAgent.__new__(DataAnalysisAgent)
|
||
agent._session_ref = None
|
||
from unittest.mock import MagicMock
|
||
agent.executor = MagicMock()
|
||
agent.executor.execute_code.return_value = {
|
||
"success": True, "output": "", "error": "",
|
||
"variables": {}, "evidence_rows": [],
|
||
"auto_exported_files": [], "prompt_saved_files": [],
|
||
}
|
||
yaml_data = {"code": "x = 1"}
|
||
result = agent._handle_generate_code("response text", yaml_data)
|
||
assert result["reasoning"] == ""
|
||
|
||
|
||
# ===========================================================================
|
||
# Task 8.3: set_session_ref
|
||
# ===========================================================================
|
||
|
||
class TestSetSessionRef:
|
||
def test_session_ref_default_none(self):
|
||
"""8.3: _session_ref defaults to None."""
|
||
agent = DataAnalysisAgent()
|
||
assert agent._session_ref is None
|
||
|
||
def test_set_session_ref(self):
|
||
"""8.3: set_session_ref stores the session reference."""
|
||
agent = DataAnalysisAgent()
|
||
|
||
class FakeSession:
|
||
rounds = []
|
||
data_files = []
|
||
|
||
session = FakeSession()
|
||
agent.set_session_ref(session)
|
||
assert agent._session_ref is session
|
||
|
||
|
||
# ===========================================================================
|
||
# Task 9.1: Prompt - intermediate data saving instructions
|
||
# ===========================================================================
|
||
|
||
class TestPromptDataSaving:
|
||
def test_data_saving_instructions_in_system_prompt(self):
|
||
"""9.1: data_analysis_system_prompt contains DATA_FILE_SAVED instructions."""
|
||
assert "[DATA_FILE_SAVED]" in data_analysis_system_prompt
|
||
assert "中间数据保存规则" in data_analysis_system_prompt
|
||
|
||
def test_data_saving_example_in_prompt(self):
|
||
"""9.1: Prompt contains example of saving and printing marker."""
|
||
assert "to_csv" in data_analysis_system_prompt
|
||
assert "session_output_dir" in data_analysis_system_prompt
|
||
|
||
|
||
# ===========================================================================
|
||
# Task 9.2: Prompt - evidence annotation instructions
|
||
# ===========================================================================
|
||
|
||
class TestPromptEvidenceAnnotation:
|
||
def test_evidence_annotation_in_report_prompt(self):
|
||
"""9.2: final_report_system_prompt contains evidence annotation instructions."""
|
||
assert "evidence:round_" in final_report_system_prompt
|
||
assert "证据标注规则" in final_report_system_prompt
|
||
|
||
def test_evidence_annotation_example(self):
|
||
"""9.2: Prompt contains example of evidence annotation."""
|
||
assert "<!-- evidence:round_3 -->" in final_report_system_prompt
|
||
|
||
|
||
# ===========================================================================
|
||
# Task 9.3: _build_final_report_prompt includes evidence
|
||
# ===========================================================================
|
||
|
||
class TestBuildFinalReportPromptEvidence:
|
||
def test_evidence_included_when_session_has_rounds(self):
|
||
"""9.3: _build_final_report_prompt includes evidence data when rounds exist."""
|
||
agent = DataAnalysisAgent.__new__(DataAnalysisAgent)
|
||
agent.analysis_results = []
|
||
agent.current_round = 2
|
||
agent.session_output_dir = "/tmp/test"
|
||
agent.data_profile = "test profile"
|
||
|
||
class FakeSession:
|
||
rounds = [
|
||
{
|
||
"round": 1,
|
||
"reasoning": "分析车型分布",
|
||
"result_summary": "执行成功，输出 DataFrame (10行×3列)",
|
||
"evidence_rows": [{"车型": "A", "数量": 42}],
|
||
},
|
||
{
|
||
"round": 2,
|
||
"reasoning": "分析模块分布",
|
||
"result_summary": "执行成功",
|
||
"evidence_rows": [],
|
||
},
|
||
]
|
||
|
||
agent._session_ref = FakeSession()
|
||
prompt = agent._build_final_report_prompt([])
|
||
assert "各轮次分析证据数据" in prompt
|
||
assert "第1轮" in prompt
|
||
assert "第2轮" in prompt
|
||
assert "车型" in prompt
|
||
|
||
def test_no_evidence_when_no_session_ref(self):
|
||
"""9.3: _build_final_report_prompt works without session ref."""
|
||
agent = DataAnalysisAgent.__new__(DataAnalysisAgent)
|
||
agent.analysis_results = []
|
||
agent.current_round = 1
|
||
agent.session_output_dir = "/tmp/test"
|
||
agent.data_profile = "test profile"
|
||
agent._session_ref = None
|
||
prompt = agent._build_final_report_prompt([])
|
||
assert "各轮次分析证据数据" not in prompt
|