test_analyzer.py

"""
测试 ResultAnalyzer 功能

这个脚本测试结果分析器的基本功能，包括统计计算和报告生成。
"""

from datetime import datetime, date
from decimal import Decimal
from nl2sql_agent import ResultAnalyzer, MockLLMClient


def test_empty_result():
    """测试空结果集"""
    print("=" * 60)
    print("测试 1: 空结果集")
    print("=" * 60)
    
    analyzer = ResultAnalyzer()
    
    result = analyzer.analyze(
        query_result=[],
        original_query="查询不存在的数据",
        sql_statement="SELECT * FROM customers WHERE id = 99999"
    )
    
    print(result)
    print("\n")


def test_numeric_data():
    """测试数值数据分析"""
    print("=" * 60)
    print("测试 2: 数值数据分析")
    print("=" * 60)
    
    analyzer = ResultAnalyzer()
    
    # 模拟订单数据
    query_result = [
        {"order_id": 1, "customer_name": "张三", "total_amount": Decimal("1500.50"), "order_date": datetime(2024, 1, 15)},
        {"order_id": 2, "customer_name": "李四", "total_amount": Decimal("2300.00"), "order_date": datetime(2024, 1, 16)},
        {"order_id": 3, "customer_name": "王五", "total_amount": Decimal("890.75"), "order_date": datetime(2024, 1, 17)},
        {"order_id": 4, "customer_name": "赵六", "total_amount": Decimal("3200.00"), "order_date": datetime(2024, 1, 18)},
        {"order_id": 5, "customer_name": "钱七", "total_amount": Decimal("1100.25"), "order_date": datetime(2024, 1, 19)},
    ]
    
    result = analyzer.analyze(
        query_result=query_result,
        original_query="查询最近的订单",
        sql_statement="SELECT order_id, customer_name, total_amount, order_date FROM orders ORDER BY order_date DESC LIMIT 5"
    )
    
    print(result)
    print("\n")


def test_text_data():
    """测试文本数据分析"""
    print("=" * 60)
    print("测试 3: 文本数据分析")
    print("=" * 60)
    
    analyzer = ResultAnalyzer()
    
    # 模拟客户数据
    query_result = [
        {"customer_id": 1, "name": "张三", "city": "北京", "country": "中国"},
        {"customer_id": 2, "name": "李四", "city": "上海", "country": "中国"},
        {"customer_id": 3, "name": "王五", "city": "北京", "country": "中国"},
        {"customer_id": 4, "name": "赵六", "city": "深圳", "country": "中国"},
        {"customer_id": 5, "name": "钱七", "city": "北京", "country": "中国"},
        {"customer_id": 6, "name": "孙八", "city": "上海", "country": "中国"},
        {"customer_id": 7, "name": "周九", "city": "广州", "country": "中国"},
        {"customer_id": 8, "name": "吴十", "city": "北京", "country": "中国"},
    ]
    
    result = analyzer.analyze(
        query_result=query_result,
        original_query="查询所有客户信息",
        sql_statement="SELECT customer_id, name, city, country FROM customers"
    )
    
    print(result)
    print("\n")


def test_with_mock_llm():
    """测试使用 Mock LLM 生成报告"""
    print("=" * 60)
    print("测试 4: 使用 Mock LLM 生成报告")
    print("=" * 60)
    
    # 创建 Mock LLM 客户端
    llm_client = MockLLMClient()
    
    # 创建带 LLM 的分析器
    analyzer = ResultAnalyzer(llm_client=llm_client)
    
    # 模拟产品销售数据
    query_result = [
        {"product_name": "iPhone 14", "sales_count": 150, "revenue": Decimal("149850.00")},
        {"product_name": "MacBook Air", "sales_count": 85, "revenue": Decimal("84915.00")},
        {"product_name": "AirPods Pro", "sales_count": 320, "revenue": Decimal("79680.00")},
        {"product_name": "iPad Pro", "sales_count": 120, "revenue": Decimal("95880.00")},
    ]
    
    result = analyzer.analyze(
        query_result=query_result,
        original_query="查询产品销售统计",
        sql_statement="""
            SELECT 
                p.product_name,
                COUNT(oi.order_item_id) as sales_count,
                SUM(oi.quantity * oi.unit_price) as revenue
            FROM products p
            JOIN order_items oi ON p.product_id = oi.product_id
            GROUP BY p.product_id, p.product_name
            ORDER BY revenue DESC
        """
    )
    
    print(result)
    print("\n")


def test_statistics_calculation():
    """测试统计信息计算"""
    print("=" * 60)
    print("测试 5: 统计信息计算")
    print("=" * 60)
    
    analyzer = ResultAnalyzer()
    
    # 模拟混合类型数据
    query_result = [
        {
            "id": 1,
            "name": "产品A",
            "price": Decimal("99.99"),
            "stock": 100,
            "category": "电子产品",
            "created_at": datetime(2024, 1, 1)
        },
        {
            "id": 2,
            "name": "产品B",
            "price": Decimal("199.99"),
            "stock": 50,
            "category": "电子产品",
            "created_at": datetime(2024, 1, 15)
        },
        {
            "id": 3,
            "name": "产品C",
            "price": Decimal("49.99"),
            "stock": 200,
            "category": "家居用品",
            "created_at": datetime(2024, 2, 1)
        },
    ]
    
    # 直接调用统计计算方法
    statistics = analyzer._calculate_statistics(query_result)
    
    print("统计信息:")
    print(f"记录数: {statistics['record_count']}")
    print(f"字段列表: {statistics['columns']}")
    print(f"字段类型: {statistics['column_types']}")
    print(f"\n数值字段统计:")
    for col, stats in statistics['numeric_stats'].items():
        print(f"  {col}:")
        print(f"    最小值: {stats['min']}")
        print(f"    最大值: {stats['max']}")
        print(f"    平均值: {stats['avg']:.2f}")
        print(f"    总和: {stats['sum']:.2f}")
    
    print(f"\n文本字段统计:")
    for col, stats in statistics['text_stats'].items():
        print(f"  {col}:")
        print(f"    唯一值数量: {stats['unique_count']}")
        if stats['most_common']:
            print(f"    最常见的值: {stats['most_common']}")
    
    print(f"\n日期字段统计:")
    for col, stats in statistics['date_stats'].items():
        print(f"  {col}:")
        print(f"    最早: {stats['earliest']}")
        print(f"    最晚: {stats['latest']}")
    
    print("\n")


if __name__ == "__main__":
    print("\n开始测试 ResultAnalyzer\n")
    
    # 运行所有测试
    test_empty_result()
    test_numeric_data()
    test_text_data()
    test_statistics_calculation()
    test_with_mock_llm()
    
    print("=" * 60)
    print("所有测试完成！")
    print("=" * 60)