#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ 测试优化后的系统功能 验证自动清理、内存处理和流式上传 """ import os import sys import time import tempfile import logging import subprocess import signal from io import BytesIO from PIL import Image import requests import json # 设置日志 logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) def start_web_server(): """启动Web服务器""" try: logger.info("🚀 启动Web服务器...") # 启动Web应用进程 process = subprocess.Popen([ sys.executable, 'web_app_vdb_production.py' ], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) logger.info(f"✅ Web服务器已启动，PID: {process.pid}") # 等待服务器启动 for i in range(10): try: response = requests.get("http://127.0.0.1:5000/", timeout=2) if response.status_code == 200: logger.info("✅ Web服务器就绪") return process except: time.sleep(2) logger.warning("⚠️ Web服务器启动超时") return process except Exception as e: logger.error(f"❌ 启动Web服务器失败: {e}") return None def stop_web_server(process): """停止Web服务器""" if process: try: process.terminate() process.wait(timeout=5) logger.info("✅ Web服务器已停止") except subprocess.TimeoutExpired: process.kill() logger.info("🔥 强制停止Web服务器") except Exception as e: logger.error(f"❌ 停止Web服务器失败: {e}") def test_optimized_file_handler(): """测试优化的文件处理器""" print("\n" + "="*60) print("测试优化的文件处理器") print("="*60) try: from optimized_file_handler import get_file_handler # 获取文件处理器实例 file_handler = get_file_handler() logger.info("✅ 文件处理器初始化成功") # 创建测试图像 test_image = Image.new('RGB', (100, 100), color='red') img_buffer = BytesIO() test_image.save(img_buffer, format='PNG') img_buffer.seek(0) # 测试文件大小判断 file_size = file_handler.get_file_size(img_buffer) is_small = file_handler.is_small_file(img_buffer) logger.info(f"测试图像大小: {file_size} bytes, 小文件: {is_small}") # 测试临时文件上下文管理器 with file_handler.temp_file_context(suffix='.png') as temp_path: logger.info(f"创建临时文件: {temp_path}") assert os.path.exists(temp_path) # 验证临时文件已被清理 assert not os.path.exists(temp_path) logger.info("✅ 临时文件自动清理成功") # 测试清理所有临时文件 file_handler.cleanup_all_temp_files() logger.info("✅ 批量清理临时文件成功") return True except Exception as e: logger.error(f"❌ 文件处理器测试失败: {e}") return False def test_memory_processing(): """测试内存处理功能""" print("\n" + "="*60) print("测试内存处理功能") print("="*60) try: from optimized_file_handler import get_file_handler file_handler = get_file_handler() # 测试文本内存处理 test_texts = [ "这是一个测试文本", "测试内存处理功能", "优化的文件处理器" ] logger.info(f"开始内存处理 {len(test_texts)} 条文本...") processed_texts = file_handler.process_text_in_memory(test_texts) if processed_texts: logger.info(f"✅ 内存处理文本成功: {len(processed_texts)} 条") for i, text_info in enumerate(processed_texts): logger.info(f" 文本 {i}: {text_info['bos_key']}") else: logger.warning("⚠️ 内存处理文本返回空结果") return len(processed_texts) > 0 except Exception as e: logger.error(f"❌ 内存处理测试失败: {e}") return False def test_web_api_optimized(): """测试优化后的Web API""" print("\n" + "="*60) print("测试优化后的Web API") print("="*60) base_url = "http://127.0.0.1:5000" try: # 测试系统初始化 logger.info("测试系统初始化...") response = requests.post(f"{base_url}/api/init") if response.status_code == 200: result = response.json() logger.info(f"✅ 系统初始化: {result.get('message')}") else: logger.warning(f"⚠️ 系统可能已初始化: {response.status_code}") # 测试文本上传（内存处理） logger.info("测试文本上传（内存处理）...") text_data = { "texts": [ "优化后的文本处理测试", "内存模式处理文本", "自动清理临时文件" ] } response = requests.post( f"{base_url}/api/upload/texts", json=text_data, headers={'Content-Type': 'application/json'} ) if response.status_code == 200: result = response.json() logger.info(f"✅ 文本上传成功: {result.get('message')}") logger.info(f" 处理方法: {result.get('processing_method')}") logger.info(f" 处理数量: {result.get('processed_texts')}") else: logger.error(f"❌ 文本上传失败: {response.status_code}") logger.error(f" 错误信息: {response.text}") # 测试图像上传（智能处理） logger.info("测试图像上传（智能处理）...") # 创建测试图像 test_image = Image.new('RGB', (200, 200), color='blue') img_buffer = BytesIO() test_image.save(img_buffer, format='PNG') img_buffer.seek(0) files = {'files': ('test_image.png', img_buffer, 'image/png')} response = requests.post(f"{base_url}/api/upload/images", files=files) if response.status_code == 200: result = response.json() logger.info(f"✅ 图像上传成功: {result.get('message')}") logger.info(f" 处理方法: {result.get('processing_methods')}") logger.info(f" 处理数量: {result.get('processed_files')}") else: logger.error(f"❌ 图像上传失败: {response.status_code}") logger.error(f" 错误信息: {response.text}") # 测试文搜文 logger.info("测试文搜文...") search_data = { "query": "优化处理", "top_k": 3 } response = requests.post( f"{base_url}/api/search/text_to_text", json=search_data, headers={'Content-Type': 'application/json'} ) if response.status_code == 200: result = response.json() logger.info(f"✅ 文搜文成功: 找到 {result.get('count')} 个结果") for i, res in enumerate(result.get('results', [])): logger.info(f" 结果 {i}: 相似度 {res['score']:.3f}") else: logger.error(f"❌ 文搜文失败: {response.status_code}") # 测试系统统计 logger.info("测试系统统计...") response = requests.get(f"{base_url}/api/stats") if response.status_code == 200: result = response.json() stats = result.get('stats', {}) logger.info(f"✅ 系统统计:") logger.info(f" 文本数量: {stats.get('text_count', 0)}") logger.info(f" 图像数量: {stats.get('image_count', 0)}") logger.info(f" 后端: {result.get('backend')}") else: logger.error(f"❌ 获取系统统计失败: {response.status_code}") return True except requests.exceptions.ConnectionError: logger.error("❌ 无法连接到Web服务器，请确保服务器正在运行") return False except Exception as e: logger.error(f"❌ Web API测试失败: {e}") return False def test_temp_file_cleanup(): """测试临时文件清理""" print("\n" + "="*60) print("测试临时文件清理") print("="*60) try: from optimized_file_handler import get_file_handler file_handler = get_file_handler() # 创建多个临时文件 temp_paths = [] for i in range(3): temp_path = file_handler.get_temp_file_for_model( BytesIO(b"test content"), f"test_{i}.txt" ) if temp_path: temp_paths.append(temp_path) logger.info(f"创建临时文件: {temp_path}") # 验证文件存在 existing_count = sum(1 for path in temp_paths if os.path.exists(path)) logger.info(f"创建的临时文件数量: {existing_count}") # 清理所有临时文件 file_handler.cleanup_all_temp_files() # 验证文件已清理 remaining_count = sum(1 for path in temp_paths if os.path.exists(path)) logger.info(f"清理后剩余文件数量: {remaining_count}") if remaining_count == 0: logger.info("✅ 临时文件清理测试成功") return True else: logger.warning(f"⚠️ 仍有 {remaining_count} 个文件未清理") return False except Exception as e: logger.error(f"❌ 临时文件清理测试失败: {e}") return False def main(): """主测试函数""" print("🚀 开始测试优化后的系统功能") print("="*60) # 启动Web服务器 web_process = start_web_server() try: test_results = [] # 运行各项测试 tests = [ ("文件处理器基础功能", test_optimized_file_handler), ("内存处理功能", test_memory_processing), ("临时文件清理", test_temp_file_cleanup), ("Web API功能", test_web_api_optimized), ] for test_name, test_func in tests: logger.info(f"\n🔍 开始测试: {test_name}") try: result = test_func() test_results.append((test_name, result)) if result: logger.info(f"✅ {test_name} - 通过") else: logger.warning(f"⚠️ {test_name} - 失败") except Exception as e: logger.error(f"❌ {test_name} - 异常: {e}") test_results.append((test_name, False)) # 输出测试总结 print("\n" + "="*60) print("测试结果总结") print("="*60) passed = sum(1 for _, result in test_results if result) total = len(test_results) for test_name, result in test_results: status = "✅ 通过" if result else "❌ 失败" print(f"{test_name}: {status}") print(f"\n总体结果: {passed}/{total} 项测试通过") if passed == total: print("🎉 所有测试通过！优化系统功能正常") else: print("⚠️ 部分测试失败，请检查相关功能") return passed == total finally: # 确保停止Web服务器 stop_web_server(web_process) if __name__ == "__main__": success = main() sys.exit(0 if success else 1)