ShopTRAINING/test/test_model_management.py

#!/usr/bin/env python
"""
测试统一模型管理功能
"""

import sys
import os
sys.path.append('server')

def test_model_manager():
    """测试模型管理器基本功能"""
    print("=== 测试模型管理器 ===")
    
    try:
        from utils.model_manager import model_manager
        
        # 测试列出模型
        models = model_manager.list_models()
        print(f"OK: 模型管理器初始化成功")
        print(f"INFO: 找到 {len(models)} 个模型")
        
        if models:
            print("\n模型列表:")
            for i, model in enumerate(models[:3]):  # 只显示前3个
                print(f"{i+1}. {model.get('filename', 'N/A')}")
                print(f"   产品: {model.get('product_id', 'N/A')}")
                print(f"   类型: {model.get('model_type', 'N/A')}")
                print(f"   模式: {model.get('training_mode', 'N/A')}")
                print(f"   版本: {model.get('version', 'N/A')}")
                if model.get('metrics'):
                    print(f"   RMSE: {model.get('metrics', {}).get('rmse', 'N/A')}")
        
        return True
        
    except Exception as e:
        print(f"FAIL: 模型管理器测试失败: {e}")
        return False

def test_api_models_endpoint():
    """测试API模型端点"""
    print("\n=== 测试API模型端点 ===")
    
    try:
        import requests
        
        # 测试获取模型列表
        response = requests.get("http://localhost:5000/api/models", timeout=5)
        
        if response.status_code == 200:
            data = response.json()
            models = data.get('data', [])
            print(f"OK: API返回 {len(models)} 个模型")
            
            if models:
                print("\nAPI返回的模型示例:")
                model = models[0]
                print(f"  模型ID: {model.get('model_id', 'N/A')}")
                print(f"  产品: {model.get('product_name', 'N/A')}")
                print(f"  类型: {model.get('model_type', 'N/A')}")
                print(f"  训练模式: {model.get('training_mode', 'N/A')}")
            
            return True
        else:
            print(f"FAIL: API返回错误状态码: {response.status_code}")
            return False
            
    except ImportError:
        print("WARN: requests模块未安装，跳过API测试")
        return True
    except Exception as e:
        # 检查是否是连接错误
        if "ConnectionError" in str(type(e)) or "连接" in str(e):
            print("WARN: 无法连接到API服务器 (可能未启动)")
            return True  # 不算失败，因为服务器可能未启动
        else:
            print(f"FAIL: API测试失败: {e}")
            return False

def test_training_and_model_save():
    """测试训练并验证模型保存"""
    print("\n=== 测试训练和模型保存 ===")
    
    try:
        from core.predictor import PharmacyPredictor
        from utils.model_manager import model_manager
        
        # 记录训练前的模型数量
        models_before = len(model_manager.list_models())
        print(f"INFO: 训练前模型数量: {models_before}")
        
        # 创建预测器并训练
        predictor = PharmacyPredictor()
        print("INFO: 开始训练测试（TCN模型，2轮次）...")
        
        metrics = predictor.train_model(
            product_id='P001',
            model_type='tcn',
            epochs=2,
            training_mode='product'
        )
        
        if metrics:
            print("OK: 训练成功完成")
            
            # 检查模型是否正确保存
            models_after = model_manager.list_models()
            new_models = len(models_after) - models_before
            
            if new_models > 0:
                print(f"OK: 新增 {new_models} 个模型")
                
                # 查找新训练的模型
                latest_models = model_manager.list_models(product_id='P001', model_type='tcn')
                if latest_models:
                    latest_model = latest_models[0]
                    print(f"INFO: 最新模型文件名: {latest_model.get('filename', 'N/A')}")
                    print(f"INFO: 模型包含管理信息: {'model_manager_info' in str(latest_model)}")
                    
                return True
            else:
                print("WARN: 训练成功但未检测到新模型文件")
                return True
        else:
            print("WARN: 训练未返回指标（可能是预期的数据不足错误）")
            return True
            
    except Exception as e:
        print(f"INFO: 训练测试异常: {e}")
        # 检查是否是预期的数据不足错误
        if "数据不足" in str(e):
            print("INFO: 这是预期的数据不足错误，错误处理正常")
            return True
        else:
            return False

def test_filename_parsing():
    """测试文件名解析功能"""
    print("\n=== 测试文件名解析 ===")
    
    try:
        from utils.model_manager import model_manager
        
        # 测试各种文件名格式
        test_filenames = [
            "tcn_product_P001_v1.pth",
            "mlstm_store_S001_P001_v2.pth", 
            "kan_global_P001_sum_v1.pth",
            "transformer_model_product_P001_v1.pth",  # 旧格式
        ]
        
        for filename in test_filenames:
            info = model_manager.parse_model_filename(filename)
            if info:
                print(f"OK: {filename}")
                print(f"    类型: {info.get('model_type')}, 产品: {info.get('product_id')}")
                print(f"    模式: {info.get('training_mode')}, 版本: {info.get('version')}")
            else:
                print(f"FAIL: 无法解析 {filename}")
        
        return True
        
    except Exception as e:
        print(f"FAIL: 文件名解析测试失败: {e}")
        return False

def main():
    """主测试函数"""
    print("开始模型管理功能测试")
    
    tests_passed = 0
    total_tests = 4
    
    # 测试模型管理器
    if test_model_manager():
        tests_passed += 1
    
    # 测试文件名解析
    if test_filename_parsing():
        tests_passed += 1
    
    # 测试API端点
    if test_api_models_endpoint():
        tests_passed += 1
    
    # 测试训练和模型保存
    if test_training_and_model_save():
        tests_passed += 1
    
    print(f"\n测试结果: {tests_passed}/{total_tests} 项测试通过")
    
    if tests_passed == total_tests:
        print("SUCCESS: 模型管理功能测试通过！")
        print("\n建议:")
        print("1. 启动API服务器: uv run ./server/api.py")
        print("2. 在前端查看模型管理页面")
        print("3. 进行一次完整训练测试模型命名")
    else:
        print("WARNING: 部分测试失败，需要检查问题")

if __name__ == "__main__":
    main()