#!/usr/bin/env python3 """ Cloudflare Memory Bank 实现测试 """ import unittest import asyncio import os import sys from unittest.mock import patch, MagicMock, AsyncMock from datetime import datetime # 添加项目根目录到Python路径 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) from src.jixia.memory.cloudflare_memory_bank import CloudflareMemoryBank, MemoryEntry class TestCloudflareMemoryBank(unittest.TestCase): """测试CloudflareMemoryBank类""" def setUp(self): """测试前的设置""" # Mock掉 aiohttp.ClientSession 以避免实际网络请求 self.patcher = patch('src.jixia.memory.cloudflare_memory_bank.aiohttp.ClientSession') self.mock_session_class = self.patcher.start() self.mock_session = AsyncMock() self.mock_session_class.return_value = self.mock_session # Mock掉 get_cloudflare_config 以避免实际读取配置 self.config_patcher = patch('src.jixia.memory.cloudflare_memory_bank.get_cloudflare_config') self.mock_get_config = self.config_patcher.start() self.mock_get_config.return_value = { 'account_id': 'test-account', 'api_token': 'test-token', 'vectorize_index': 'test-index', 'embed_model': '@cf/baai/bge-m3', 'autorag_domain': 'test.example.com' } # 创建CloudflareMemoryBank实例 self.memory_bank = CloudflareMemoryBank() # 重置一些内部状态 self.memory_bank.config = self.mock_get_config.return_value self.memory_bank.account_id = 'test-account' self.memory_bank.api_token = 'test-token' self.memory_bank.vectorize_index = 'test-index' self.memory_bank.embed_model = '@cf/baai/bge-m3' self.memory_bank.autorag_domain = 'test.example.com' def tearDown(self): """测试后的清理""" self.patcher.stop() self.config_patcher.stop() def test_init(self): """测试初始化""" self.assertEqual(self.memory_bank.account_id, "test-account") self.assertEqual(self.memory_bank.api_token, "test-token") self.assertEqual(self.memory_bank.vectorize_index, "test-index") self.assertEqual(self.memory_bank.embed_model, "@cf/baai/bge-m3") self.assertEqual(self.memory_bank.autorag_domain, "test.example.com") async def test_create_memory_bank(self): """测试创建记忆空间""" memory_bank_id = await self.memory_bank.create_memory_bank("tieguaili") # 验证返回的ID格式 self.assertEqual(memory_bank_id, "cf_memory_tieguaili") async def test_create_memory_bank_with_display_name(self): """测试创建记忆空间时指定显示名称""" memory_bank_id = await self.memory_bank.create_memory_bank( "tieguaili", "铁拐李的专属记忆银行" ) # 验证返回的ID格式 self.assertEqual(memory_bank_id, "cf_memory_tieguaili") async def test_generate_embedding(self): """测试生成嵌入向量""" # Mock响应 mock_response = AsyncMock() mock_response.status = 200 mock_response.json = AsyncMock(return_value={ "result": { "data": [ { "embedding": [0.1, 0.2, 0.3, 0.4, 0.5] } ] } }) # Mock session.post self.mock_session.post.return_value.__aenter__.return_value = mock_response # 调用方法 embedding = await self.memory_bank._generate_embedding("测试文本") # 验证结果 self.assertEqual(embedding, [0.1, 0.2, 0.3, 0.4, 0.5]) # 验证调用了正确的URL和参数 expected_url = "https://api.cloudflare.com/client/v4/accounts/test-account/ai/run/@cf/baai/bge-m3" self.mock_session.post.assert_called_once() call_args = self.mock_session.post.call_args self.assertEqual(call_args[0][0], expected_url) self.assertEqual(call_args[1]['json'], {"text": ["测试文本"]}) async def test_generate_embedding_api_error(self): """测试生成嵌入向量时API错误""" # Mock响应 mock_response = AsyncMock() mock_response.status = 500 mock_response.text = AsyncMock(return_value="Internal Server Error") # Mock session.post self.mock_session.post.return_value.__aenter__.return_value = mock_response # 验证抛出异常 with self.assertRaises(Exception) as context: await self.memory_bank._generate_embedding("测试文本") self.assertIn("Failed to generate embedding", str(context.exception)) async def test_add_memory(self): """测试添加记忆""" # Mock _generate_embedding 方法 with patch.object(self.memory_bank, '_generate_embedding', new=AsyncMock(return_value=[0.1, 0.2, 0.3])) as mock_embed: # Mock upsert 响应 mock_response = AsyncMock() mock_response.status = 200 mock_response.json = AsyncMock(return_value={"result": {"upserted": 1}}) # Mock session.post self.mock_session.post.return_value.__aenter__.return_value = mock_response # 添加记忆 memory_id = await self.memory_bank.add_memory( agent_name="tieguaili", content="在讨论NVIDIA股票时，我倾向于逆向思维，关注潜在风险。", memory_type="preference", debate_topic="NVIDIA投资分析", metadata={"source": "manual"} ) # 验证返回的ID格式 (以mem_开头) self.assertTrue(memory_id.startswith("mem_tieguaili_")) # 验证调用了生成嵌入的方法 mock_embed.assert_called_once_with("在讨论NVIDIA股票时，我倾向于逆向思维，关注潜在风险。") # 验证调用了upsert API self.mock_session.post.assert_called() # 验证upsert调用的参数 upsert_call = None for call in self.mock_session.post.call_args_list: if 'vectorize/indexes/test-index/upsert' in call[0][0]: upsert_call = call break self.assertIsNotNone(upsert_call) call_args, call_kwargs = upsert_call self.assertIn("vectorize/indexes/test-index/upsert", call_args[0]) self.assertIn("vectors", call_kwargs['json']) async def test_add_memory_api_error(self): """测试添加记忆时API错误""" # Mock _generate_embedding 方法 with patch.object(self.memory_bank, '_generate_embedding', new=AsyncMock(return_value=[0.1, 0.2, 0.3])): # Mock upsert 响应 mock_response = AsyncMock() mock_response.status = 500 mock_response.text = AsyncMock(return_value="Internal Server Error") # Mock session.post self.mock_session.post.return_value.__aenter__.return_value = mock_response # 验证抛出异常 with self.assertRaises(Exception) as context: await self.memory_bank.add_memory( agent_name="tieguaili", content="测试内容" ) self.assertIn("Failed to upsert memory", str(context.exception)) async def test_search_memories(self): """测试搜索记忆""" # Mock _generate_embedding 方法 with patch.object(self.memory_bank, '_generate_embedding', new=AsyncMock(return_value=[0.1, 0.2, 0.3])) as mock_embed: # Mock query 响应 mock_response = AsyncMock() mock_response.status = 200 mock_response.json = AsyncMock(return_value={ "result": { "matches": [ { "metadata": { "content": "在讨论NVIDIA股票时，我倾向于逆向思维，关注潜在风险。", "memory_type": "preference", "agent_name": "tieguaili" }, "score": 0.95 } ] } }) # Mock session.post self.mock_session.post.return_value.__aenter__.return_value = mock_response # 搜索记忆 results = await self.memory_bank.search_memories( agent_name="tieguaili", query="NVIDIA", memory_type="preference", limit=5 ) # 验证结果 self.assertEqual(len(results), 1) self.assertEqual(results[0]["content"], "在讨论NVIDIA股票时，我倾向于逆向思维，关注潜在风险。") self.assertEqual(results[0]["relevance_score"], 0.95) # 验证调用了生成嵌入的方法 mock_embed.assert_called_once_with("NVIDIA") # 验证调用了query API self.mock_session.post.assert_called() # 验证query调用的参数 query_call = None for call in self.mock_session.post.call_args_list: if 'vectorize/indexes/test-index/query' in call[0][0]: query_call = call break self.assertIsNotNone(query_call) call_args, call_kwargs = query_call self.assertIn("vectorize/indexes/test-index/query", call_args[0]) self.assertIn("vector", call_kwargs['json']) self.assertIn("filter", call_kwargs['json']) self.assertEqual(call_kwargs['json']['filter'], {"agent_name": "tieguaili", "memory_type": "preference"}) async def test_search_memories_api_error(self): """测试搜索记忆时API错误""" # Mock _generate_embedding 方法 with patch.object(self.memory_bank, '_generate_embedding', new=AsyncMock(return_value=[0.1, 0.2, 0.3])): # Mock query 响应 mock_response = AsyncMock() mock_response.status = 500 mock_response.text = AsyncMock(return_value="Internal Server Error") # Mock session.post self.mock_session.post.return_value.__aenter__.return_value = mock_response # 验证返回空列表而不是抛出异常 results = await self.memory_bank.search_memories( agent_name="tieguaili", query="NVIDIA" ) self.assertEqual(results, []) async def test_get_agent_context(self): """测试获取智能体上下文""" # Mock search_memories 方法 with patch.object(self.memory_bank, 'search_memories', new=AsyncMock()) as mock_search: # 设置mock返回值 mock_search.side_effect = [ [ # conversation memories {"content": "NVIDIA的估值过高，存在泡沫风险。", "relevance_score": 0.9} ], [ # preference memories {"content": "倾向于逆向思维，关注潜在风险。", "relevance_score": 0.8} ], [ # strategy memories {"content": "使用技术分析策略。", "relevance_score": 0.7} ] ] # 获取上下文 context = await self.memory_bank.get_agent_context("tieguaili", "NVIDIA投资分析") # 验证上下文包含预期内容 self.assertIn("# 铁拐李的记忆上下文", context) self.assertIn("## 历史对话记忆", context) self.assertIn("## 偏好记忆", context) self.assertIn("## 策略记忆", context) self.assertIn("NVIDIA的估值过高，存在泡沫风险。", context) self.assertIn("倾向于逆向思维，关注潜在风险。", context) self.assertIn("使用技术分析策略。", context) async def test_get_agent_context_no_memories(self): """测试获取智能体上下文但无相关记忆""" # Mock search_memories 方法 with patch.object(self.memory_bank, 'search_memories', new=AsyncMock(return_value=[])): # 获取上下文 context = await self.memory_bank.get_agent_context("tieguaili", "NVIDIA投资分析") # 验证上下文包含暂无相关记忆的提示 self.assertIn("# 铁拐李的记忆上下文", context) self.assertIn("暂无相关记忆。", context) async def test_save_debate_session(self): """测试保存辩论会话""" # Mock add_memory 方法 with patch.object(self.memory_bank, 'add_memory', new=AsyncMock()) as mock_add: conversation_history = [ {"agent": "tieguaili", "content": "NVIDIA的估值过高，存在泡沫风险。"}, {"agent": "lvdongbin", "content": "NVIDIA在AI领域的领先地位不可忽视。"}, {"agent": "tieguaili", "content": "但我们需要考虑竞争加剧和增长放缓的可能性。"} ] outcomes = { "winner": "lvdongbin", "insights": { "tieguaili": "铁拐李的风险意识值得肯定，但在AI趋势的判断上略显保守。" } } # 保存辩论会话 await self.memory_bank.save_debate_session( debate_topic="NVIDIA投资分析", participants=["tieguaili", "lvdongbin"], conversation_history=conversation_history, outcomes=outcomes ) # 验证调用了add_memory两次（对话总结和策略洞察） self.assertEqual(mock_add.call_count, 2) # 验证第一次调用是对话总结 call_args1 = mock_add.call_args_list[0][1] self.assertEqual(call_args1['agent_name'], 'tieguaili') self.assertEqual(call_args1['memory_type'], 'conversation') self.assertIn('铁拐李在本次辩论中的主要观点', call_args1['content']) # 验证第二次调用是策略洞察 call_args2 = mock_add.call_args_list[1][1] self.assertEqual(call_args2['agent_name'], 'tieguaili') self.assertEqual(call_args2['memory_type'], 'strategy') self.assertIn('铁拐李的风险意识值得肯定', call_args2['content']) def test_summarize_conversation(self): """测试对话总结""" conversation_history = [ {"agent": "tieguaili", "content": "第一点看法：NVIDIA的估值过高，存在泡沫风险。"}, {"agent": "lvdongbin", "content": "NVIDIA在AI领域的领先地位不可忽视。"}, {"agent": "tieguaili", "content": "第二点看法：我们需要考虑竞争加剧和增长放缓的可能性。"}, {"agent": "tieguaili", "content": "第三点看法：从技术分析角度看，股价已出现超买信号。"} ] summary = self.memory_bank._summarize_conversation(conversation_history, "tieguaili") # 验证总结包含预期内容 self.assertIn("铁拐李在本次辩论中的主要观点", summary) self.assertIn("第一点看法：NVIDIA的估值过高，存在泡沫风险。", summary) self.assertIn("第二点看法：我们需要考虑竞争加剧和增长放缓的可能性。", summary) self.assertIn("第三点看法：从技术分析角度看，股价已出现超买信号。", summary) def test_extract_strategy_insight_winner(self): """测试提取策略洞察 - 获胜者""" outcomes = { "winner": "tieguaili", "insights": {} } insight = self.memory_bank._extract_strategy_insight(outcomes, "tieguaili") self.assertIn("铁拐李在本次辩论中获胜", insight) def test_extract_strategy_insight_from_insights(self): """测试从洞察中提取策略洞察""" outcomes = { "winner": "lvdongbin", "insights": { "tieguaili": "铁拐李的风险意识值得肯定，但在AI趋势的判断上略显保守。" } } insight = self.memory_bank._extract_strategy_insight(outcomes, "tieguaili") self.assertEqual(insight, "铁拐李的风险意识值得肯定，但在AI趋势的判断上略显保守。") if __name__ == '__main__': # 创建一个异步测试运行器 def run_async_test(test_case): """运行异步测试用例""" loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) try: return loop.run_until_complete(test_case) finally: loop.close() # 获取所有以test_开头的异步方法并运行它们 suite = unittest.TestSuite() test_instance = TestCloudflareMemoryBank() test_instance.setUp() test_instance.addCleanup(test_instance.tearDown) # 添加同步测试 suite.addTest(TestCloudflareMemoryBank('test_init')) suite.addTest(TestCloudflareMemoryBank('test_summarize_conversation')) suite.addTest(TestCloudflareMemoryBank('test_extract_strategy_insight_winner')) suite.addTest(TestCloudflareMemoryBank('test_extract_strategy_insight_from_insights')) # 添加异步测试 async_tests = [ 'test_create_memory_bank', 'test_create_memory_bank_with_display_name', 'test_generate_embedding', 'test_generate_embedding_api_error', 'test_add_memory', 'test_add_memory_api_error', 'test_search_memories', 'test_search_memories_api_error', 'test_get_agent_context', 'test_get_agent_context_no_memories', 'test_save_debate_session' ] for test_name in async_tests: test_method = getattr(test_instance, test_name) suite.addTest(unittest.FunctionTestCase(lambda tm=test_method: run_async_test(tm()))) # 运行测试 runner = unittest.TextTestRunner(verbosity=2) runner.run(suite)