RagFlow_Plugs/test/test_mcp_hybrid_search.py

#!/usr/bin/env python3
"""
测试MCP服务的混合检索功能
"""

import sys
import os
import json
import unittest
import requests
from typing import Dict, Any

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

class TestMCPHybridSearch(unittest.TestCase):
    """测试MCP服务的混合检索功能"""
    
    def setUp(self):
        """设置测试环境"""
        # MCP服务的基础URL
        self.base_url = "http://localhost:18000"
        
        # 测试数据
        # 测试图片
        self.test_image = r"D:\project\work\ragflow_plugs\book\output\temp\2.png"
        # 测试文本查询
        self.test_text_query = "卡梅拉"

    
    def test_hybrid_search(self):
        """测试混合检索API"""
        print("测试混合检索API...")
        
        # 构建请求数据
        data = {
            "text_query": self.test_text_query,
            "image": self.test_image,
            "topn": 2
        }
        
        # 发送请求
        response = requests.post(
            f"{self.base_url}/tools/hybrid_search",
            json=data
        )
        
        # 验证响应
        self.assertEqual(response.status_code, 200, f"请求失败: {response.text}")
        result = response.json()
        self.assertTrue(result["success"], f"API调用失败: {result.get('message', '未知错误')}")
        self.assertIn("output", result, "响应中缺少hits字段")
        self.assertIn("total", result, "响应中缺少total字段")
        self.assertIsInstance(result["output"], list, "hits字段应该是一个列表")
        self.assertIsInstance(result["total"], int, "total字段应该是一个整数")
        
        print(f"✓ 混合检索API测试通过，命中数量: {result['total']}")
    
 

if __name__ == "__main__":
    """运行测试"""
    print("开始测试MCP服务的混合检索功能...\n")
    
    # 创建测试套件
    suite = unittest.TestLoader().loadTestsFromTestCase(TestMCPHybridSearch)
    
    # 运行测试
    runner = unittest.TextTestRunner(verbosity=2)
    result = runner.run(suite)
    
    print(f"\n测试完成，共运行 {result.testsRun} 个测试，成功 {result.testsRun - len(result.failures) - len(result.errors)} 个，失败 {len(result.failures)} 个，错误 {len(result.errors)} 个")