# -*- 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 "" 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