Excel MCP Server

"""
Excel MCP Server 测试用例
"""

import json
import tempfile
import pandas as pd
from pathlib import Path
import sys
import os

# 添加src目录到Python路径
src_path = Path(__file__).parent.parent / "src"
sys.path.insert(0, str(src_path))

from excel_mcp.config import ConfigManager
from excel_mcp.security import SecurityValidator
from excel_mcp.file_manager import FileManager
from excel_mcp.data_processor import DataProcessor
from excel_mcp.sheet_operations import SheetOperations


def test_basic_functionality():
    """测试基本功能"""
    print("开始测试Excel MCP Server基本功能...")
    
    # 初始化组件
    config_manager = ConfigManager()
    security_validator = SecurityValidator(config_manager)
    file_manager = FileManager(config_manager, security_validator)
    data_processor = DataProcessor(file_manager)
    sheet_operations = SheetOperations(file_manager)
    
    # 创建测试数据目录
    test_data_dir = Path("data")
    test_data_dir.mkdir(exist_ok=True)
    
    # 测试文件路径
    test_file = test_data_dir / "test_excel.xlsx"
    
    print(f"测试文件路径: {test_file}")
    
    # 1. 测试创建Excel文件
    print("\n1. 测试创建Excel文件...")
    success = file_manager.create_workbook(str(test_file), ["Sales", "Products"])
    if success:
        print("✓ Excel文件创建成功")
    else:
        print("✗ Excel文件创建失败")
        return False
    
    # 2. 测试获取文件信息
    print("\n2. 测试获取文件信息...")
    info = file_manager.get_file_info(str(test_file))
    if info:
        print(f"✓ 文件信息获取成功: {info['sheet_names']}")
    else:
        print("✗ 文件信息获取失败")
        return False
    
    # 3. 测试写入数据
    print("\n3. 测试写入数据...")
    test_data = [
        ["Product", "Price", "Quantity"],
        ["Apple", 1.50, 100],
        ["Banana", 0.80, 150],
        ["Orange", 2.00, 80]
    ]
    
    result = data_processor.write_range(str(test_file), "Sales", "A1", test_data)
    if result.get("success"):
        print("✓ 数据写入成功")
    else:
        print(f"✗ 数据写入失败: {result.get('error')}")
        return False
    
    # 4. 测试读取数据
    print("\n4. 测试读取数据...")
    result = data_processor.read_range(str(test_file), "Sales", "A1:C4")
    if result.get("success"):
        print(f"✓ 数据读取成功: {len(result['data'])} 行数据")
        print(f"数据内容: {result['data']}")
    else:
        print(f"✗ 数据读取失败: {result.get('error')}")
        return False
    
    # 5. 测试查询数据
    print("\n5. 测试查询数据...")
    result = data_processor.query_data(str(test_file), "Sales")
    if result.get("success"):
        print(f"✓ 数据查询成功: {result['row_count']} 行数据")
    else:
        print(f"✗ 数据查询失败: {result.get('error')}")
        return False
    
    # 6. 测试工作表操作
    print("\n6. 测试工作表操作...")
    
    # 创建新工作表
    result = sheet_operations.create_sheet(str(test_file), "Inventory")
    if result.get("success"):
        print("✓ 新工作表创建成功")
    else:
        print(f"✗ 新工作表创建失败: {result.get('error')}")
        return False
    
    # 列出所有工作表
    result = sheet_operations.list_sheets(str(test_file))
    if result.get("success"):
        print(f"✓ 工作表列表获取成功: {result['sheet_names']}")
    else:
        print(f"✗ 工作表列表获取失败: {result.get('error')}")
        return False
    
    # 7. 测试公式计算
    print("\n7. 测试公式计算...")
    result = data_processor.calculate_formula(str(test_file), "Sales", "D2", "B2*C2")
    if result.get("success"):
        print(f"✓ 公式计算成功: {result['value']}")
    else:
        print(f"✗ 公式计算失败: {result.get('error')}")
        return False
    
    print("\n🎉 所有基本功能测试通过！")
    return True


def test_mcp_tools():
    """测试MCP工具"""
    print("\n开始测试MCP工具...")
    
    # 这里我们模拟MCP工具调用
    test_data_dir = Path("data")
    test_file = test_data_dir / "mcp_test.xlsx"
    
    # 导入MCP工具函数
    from excel_mcp.server import (
        initialize_components,
        create_excel_file,
        write_excel_file,
        read_excel_file,
        get_excel_info,
        list_worksheets
    )
    
    # 初始化组件
    initialize_components()
    
    print("\n1. 测试创建Excel文件MCP工具...")
    result = create_excel_file(str(test_file), ["Data", "Summary"])
    if result.get("success"):
        print("✓ MCP创建文件工具测试成功")
    else:
        print(f"✗ MCP创建文件工具测试失败: {result.get('error')}")
        return False
    
    print("\n2. 测试写入数据MCP工具...")
    test_data = [
        ["Name", "Age", "City"],
        ["Alice", 25, "New York"],
        ["Bob", 30, "London"],
        ["Charlie", 35, "Tokyo"]
    ]
    result = write_excel_file(str(test_file), test_data, "Data", "A1")
    if result.get("success"):
        print("✓ MCP写入数据工具测试成功")
    else:
        print(f"✗ MCP写入数据工具测试失败: {result.get('error')}")
        return False
    
    print("\n3. 测试读取文件MCP工具...")
    result = read_excel_file(str(test_file), "Data")
    if result.get("success"):
        print(f"✓ MCP读取文件工具测试成功: 读取了 {len(result['data'])} 行数据")
    else:
        print(f"✗ MCP读取文件工具测试失败: {result.get('error')}")
        return False
    
    print("\n4. 测试获取文件信息MCP工具...")
    result = get_excel_info(str(test_file))
    if result.get("success"):
        print(f"✓ MCP获取文件信息工具测试成功: {result['sheet_names']}")
    else:
        print(f"✗ MCP获取文件信息工具测试失败: {result.get('error')}")
        return False
    
    print("\n🎉 所有MCP工具测试通过！")
    return True


def create_sample_data():
    """创建示例数据文件"""
    print("\n创建示例数据文件...")
    
    test_data_dir = Path("data")
    test_data_dir.mkdir(exist_ok=True)
    
    # 创建销售数据示例
    sales_data = {
        'Product': ['Apple', 'Banana', 'Orange', 'Grape', 'Strawberry'],
        'Price': [1.50, 0.80, 2.00, 3.50, 4.00],
        'Quantity': [100, 150, 80, 60, 40],
        'Total': [150.0, 120.0, 160.0, 210.0, 160.0]
    }
    
    df = pd.DataFrame(sales_data)
    sales_file = test_data_dir / "sales_example.xlsx"
    df.to_excel(sales_file, sheet_name="Sales", index=False)
    print(f"✓ 创建销售数据示例: {sales_file}")
    
    # 创建员工数据示例
    employee_data = {
        'Name': ['Alice Johnson', 'Bob Smith', 'Charlie Brown', 'Diana Prince', 'Eve Davis'],
        'Department': ['Sales', 'Marketing', 'IT', 'HR', 'Finance'],
        'Salary': [50000, 55000, 65000, 60000, 58000],
        'Years': [3, 5, 2, 7, 4]
    }
    
    df = pd.DataFrame(employee_data)
    employee_file = test_data_dir / "employee_example.xlsx"
    df.to_excel(employee_file, sheet_name="Employees", index=False)
    print(f"✓ 创建员工数据示例: {employee_file}")
    
    return True


def main():
    """主测试函数"""
    print("=== Excel MCP Server 测试套件 ===")
    
    try:
        # 创建示例数据
        create_sample_data()
        
        # 测试基本功能
        if not test_basic_functionality():
            print("\n❌ 基本功能测试失败")
            return False
        
        # 测试MCP工具
        if not test_mcp_tools():
            print("\n❌ MCP工具测试失败")
            return False
        
        print("\n🎊 所有测试都通过了！Excel MCP Server 已准备就绪。")
        print("\n接下来的步骤：")
        print("1. 按照 INSTALL.md 中的说明配置Claude Desktop")
        print("2. 重启Claude Desktop")
        print("3. 在Claude中测试Excel功能")
        
        return True
        
    except Exception as e:
        print(f"\n❌ 测试过程中发生错误: {e}")
        import traceback
        traceback.print_exc()
        return False


if __name__ == "__main__":
    main()