import os import sys import shutil import json # 将项目根目录添加到系统路径，以便导入server模块 project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..')) sys.path.insert(0, project_root) from server.utils.file_save import ModelPathManager def run_tests(): """执行所有路径生成逻辑的测试""" # --- 测试设置 --- test_base_dir = 'test_saved_models' if os.path.exists(test_base_dir): shutil.rmtree(test_base_dir) # 清理旧的测试目录 path_manager = ModelPathManager(base_dir=test_base_dir) model_type = 'mlstm' print("="*50) print("🚀 开始测试 ModelPathManager 路径生成逻辑...") print(f"测试根目录: {os.path.abspath(test_base_dir)}") print("="*50) # --- 1. 按店铺训练 (Store Training) 测试 --- print("\n--- 🧪 1. 按店铺训练 (Store Training) ---") # a) 店铺训练 - 所有药品 print("\n[1a] 场景: 店铺训练 - 所有药品") store_payload_all = { 'store_id': 'S001', 'model_type': model_type, 'training_mode': 'store', 'product_scope': 'all' } payload = store_payload_all.copy() payload.pop('model_type', None) payload.pop('training_mode', None) paths_store_all = path_manager.get_model_paths(training_mode='store', model_type=model_type, **payload) print(f" - Identifier: {paths_store_all['identifier']}") print(f" - Version Dir: {paths_store_all['version_dir']}") assert f"store_S001_products_all_{model_type}" == paths_store_all['identifier'] expected_path = os.path.join(test_base_dir, 'store', 'S001_all', model_type, 'v1') assert os.path.normpath(expected_path) == os.path.normpath(paths_store_all['version_dir']) # b) 店铺训练 - 特定药品 (使用哈希) print("\n[1b] 场景: 店铺训练 - 特定药品 (使用哈希)") store_payload_specific = { 'store_id': 'S002', 'model_type': model_type, 'training_mode': 'store', 'product_scope': 'specific', 'product_ids': ['P001', 'P005', 'P002'] } payload = store_payload_specific.copy() payload.pop('model_type', None) payload.pop('training_mode', None) paths_store_specific = path_manager.get_model_paths(training_mode='store', model_type=model_type, **payload) hashed_ids = path_manager._hash_ids(['P001', 'P005', 'P002']) print(f" - Hashed IDs: {hashed_ids}") print(f" - Identifier: {paths_store_specific['identifier']}") print(f" - Version Dir: {paths_store_specific['version_dir']}") assert f"store_S002_products_{hashed_ids}_{model_type}" == paths_store_specific['identifier'] expected_path = os.path.join(test_base_dir, 'store', f'S002_{hashed_ids}', model_type, 'v1') assert os.path.normpath(expected_path) == os.path.normpath(paths_store_specific['version_dir']) # c) 店铺训练 - 单个指定药品 print("\n[1c] 场景: 店铺训练 - 单个指定药品") store_payload_single_product = { 'store_id': 'S003', 'model_type': model_type, 'training_mode': 'store', 'product_scope': 'specific', 'product_ids': ['P789'] } payload = store_payload_single_product.copy() payload.pop('model_type', None) payload.pop('training_mode', None) paths_store_single_product = path_manager.get_model_paths(training_mode='store', model_type=model_type, **payload) print(f" - Identifier: {paths_store_single_product['identifier']}") print(f" - Version Dir: {paths_store_single_product['version_dir']}") assert f"store_S003_products_P789_{model_type}" == paths_store_single_product['identifier'] expected_path = os.path.join(test_base_dir, 'store', 'S003_P789', model_type, 'v1') assert os.path.normpath(expected_path) == os.path.normpath(paths_store_single_product['version_dir']) # --- 2. 按药品训练 (Product Training) 测试 --- print("\n--- 🧪 2. 按药品训练 (Product Training) ---") # a) 药品训练 - 所有店铺 print("\n[2a] 场景: 药品训练 - 所有店铺") product_payload_all = { 'product_id': 'P123', 'model_type': model_type, 'training_mode': 'product', 'store_id': None # 明确测试 None 的情况 } payload = product_payload_all.copy() payload.pop('model_type', None) payload.pop('training_mode', None) paths_product_all = path_manager.get_model_paths(training_mode='product', model_type=model_type, **payload) print(f" - Identifier: {paths_product_all['identifier']}") print(f" - Version Dir: {paths_product_all['version_dir']}") assert f"product_P123_scope_all_{model_type}" == paths_product_all['identifier'] expected_path = os.path.join(test_base_dir, 'product', 'P123_all', model_type, 'v1') assert os.path.normpath(expected_path) == os.path.normpath(paths_product_all['version_dir']) # b) 药品训练 - 特定店铺 print("\n[2b] 场景: 药品训练 - 特定店铺") product_payload_specific = { 'product_id': 'P456', 'store_id': 'S003', 'model_type': model_type, 'training_mode': 'product' } payload = product_payload_specific.copy() payload.pop('model_type', None) payload.pop('training_mode', None) paths_product_specific = path_manager.get_model_paths(training_mode='product', model_type=model_type, **payload) print(f" - Identifier: {paths_product_specific['identifier']}") print(f" - Version Dir: {paths_product_specific['version_dir']}") assert f"product_P456_scope_S003_{model_type}" == paths_product_specific['identifier'] expected_path = os.path.join(test_base_dir, 'product', 'P456_S003', model_type, 'v1') assert os.path.normpath(expected_path) == os.path.normpath(paths_product_specific['version_dir']) # --- 3. 全局训练 (Global Training) 测试 --- print("\n--- 🧪 3. 全局训练 (Global Training) ---") # a) 全局训练 - 所有数据 print("\n[3a] 场景: 全局训练 - 所有数据") global_payload_all = { 'model_type': model_type, 'training_mode': 'global', 'training_scope': 'all', 'aggregation_method': 'sum' } payload = global_payload_all.copy() payload.pop('model_type', None) payload.pop('training_mode', None) paths_global_all = path_manager.get_model_paths(training_mode='global', model_type=model_type, **payload) print(f" - Identifier: {paths_global_all['identifier']}") print(f" - Version Dir: {paths_global_all['version_dir']}") assert f"global_all_agg_sum_{model_type}" == paths_global_all['identifier'] expected_path = os.path.join(test_base_dir, 'global', 'all', 'sum', model_type, 'v1') assert os.path.normpath(expected_path) == os.path.normpath(paths_global_all['version_dir']) # a2) 全局训练 - 所有数据 (使用 all_stores_all_products) print("\n[3a2] 场景: 全局训练 - 所有数据 (使用 'all_stores_all_products')") global_payload_all_alt = { 'model_type': model_type, 'training_mode': 'global', 'training_scope': 'all_stores_all_products', 'aggregation_method': 'sum' } payload = global_payload_all_alt.copy() payload.pop('model_type', None) payload.pop('training_mode', None) paths_global_all_alt = path_manager.get_model_paths(training_mode='global', model_type=model_type, **payload) assert f"global_all_agg_sum_{model_type}" == paths_global_all_alt['identifier'] assert os.path.normpath(expected_path) == os.path.normpath(paths_global_all_alt['version_dir']) # b) 全局训练 - 自定义范围 (使用哈希) print("\n[3b] 场景: 全局训练 - 自定义范围 (使用哈希)") global_payload_custom = { 'model_type': model_type, 'training_mode': 'global', 'training_scope': 'custom', 'aggregation_method': 'mean', 'store_ids': ['S001', 'S003'], 'product_ids': ['P001', 'P002'] } payload = global_payload_custom.copy() payload.pop('model_type', None) payload.pop('training_mode', None) paths_global_custom = path_manager.get_model_paths(training_mode='global', model_type=model_type, **payload) s_hash = path_manager._hash_ids(['S001', 'S003']) p_hash = path_manager._hash_ids(['P001', 'P002']) print(f" - Store Hash: {s_hash}, Product Hash: {p_hash}") print(f" - Identifier: {paths_global_custom['identifier']}") print(f" - Version Dir: {paths_global_custom['version_dir']}") assert f"global_custom_s_{s_hash}_p_{p_hash}_agg_mean_{model_type}" == paths_global_custom['identifier'] expected_path = os.path.join(test_base_dir, 'global', 'custom', s_hash, p_hash, 'mean', model_type, 'v1') assert os.path.normpath(expected_path) == os.path.normpath(paths_global_custom['version_dir']) # c) 全局训练 - 单个店铺 print("\n[3c] 场景: 全局训练 - 单个店铺") global_payload_single_store = { 'model_type': model_type, 'training_mode': 'global', 'training_scope': 'selected_stores', 'aggregation_method': 'mean', 'store_ids': ['S007'] } payload = global_payload_single_store.copy() payload.pop('model_type', None) payload.pop('training_mode', None) paths_global_single_store = path_manager.get_model_paths(training_mode='global', model_type=model_type, **payload) print(f" - Identifier: {paths_global_single_store['identifier']}") print(f" - Version Dir: {paths_global_single_store['version_dir']}") assert f"global_stores_S007_agg_mean_{model_type}" == paths_global_single_store['identifier'] expected_path = os.path.join(test_base_dir, 'global', 'stores', 'S007', 'mean', model_type, 'v1') assert os.path.normpath(expected_path) == os.path.normpath(paths_global_single_store['version_dir']) # d) 全局训练 - 自定义范围 (单ID) print("\n[3d] 场景: 全局训练 - 自定义范围 (单ID)") global_payload_custom_single = { 'model_type': model_type, 'training_mode': 'global', 'training_scope': 'custom', 'aggregation_method': 'mean', 'store_ids': ['S008'], 'product_ids': ['P888'] } payload = global_payload_custom_single.copy() payload.pop('model_type', None) payload.pop('training_mode', None) paths_global_custom_single = path_manager.get_model_paths(training_mode='global', model_type=model_type, **payload) print(f" - Identifier: {paths_global_custom_single['identifier']}") print(f" - Version Dir: {paths_global_custom_single['version_dir']}") assert f"global_custom_s_S008_p_P888_agg_mean_{model_type}" == paths_global_custom_single['identifier'] expected_path = os.path.join(test_base_dir, 'global', 'custom', 'S008', 'P888', 'mean', model_type, 'v1') assert os.path.normpath(expected_path) == os.path.normpath(paths_global_custom_single['version_dir']) # --- 4. 版本管理测试 --- print("\n--- 🧪 4. 版本管理测试 ---") print("\n[4a] 场景: 多次调用同一训练，版本号递增") # 第一次训练 path_manager.save_version_info(paths_store_all['identifier'], paths_store_all['version']) print(f" - 保存版本: {paths_store_all['identifier']} -> v{paths_store_all['version']}") # 第二次训练 payload = store_payload_all.copy() payload.pop('model_type', None) payload.pop('training_mode', None) paths_store_all_v2 = path_manager.get_model_paths(training_mode='store', model_type=model_type, **payload) print(f" - 获取新版本: {paths_store_all_v2['identifier']} -> v{paths_store_all_v2['version']}") assert paths_store_all_v2['version'] == 2 expected_path = os.path.join(test_base_dir, 'store', 'S001_all', model_type, 'v2') assert os.path.normpath(expected_path) == os.path.normpath(paths_store_all_v2['version_dir']) # 验证 versions.json 文件 with open(path_manager.versions_file, 'r') as f: versions_data = json.load(f) print(f" - versions.json 内容: {versions_data}") assert versions_data[paths_store_all['identifier']] == 1 print("\n="*50) print("✅ 所有测试用例通过！") print("="*50) # --- 清理 --- shutil.rmtree(test_base_dir) print(f"🗑️ 测试目录 '{test_base_dir}' 已清理。") if __name__ == '__main__': run_tests()