"""HeadroomCompressor 单元测试 所有测试使用 mock headroom 模块,无需安装 headroom-ai。 """ import time from collections import OrderedDict from unittest.mock import MagicMock, patch import pytest from agentkit.core.headroom_compressor import ( HeadroomCompressor, _is_code_content, _is_json_content, ) # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def _make_headroom_compress_mock(return_content="compressed"): """创建 mock headroom.compress 函数,返回带有 messages 属性的结果对象""" mock_result = MagicMock() mock_result.messages = [{"role": "user", "content": return_content}] return mock_result def _long_json_content(): """生成超过 min_length 的 JSON 内容""" import json items = [{"id": i, "name": f"item_{i}", "description": f"description for item {i}"} for i in range(50)] return json.dumps({"items": items}) def _long_code_content(): """生成超过 min_length 的代码内容""" lines = [] for i in range(50): lines.append(f"def function_{i}():") lines.append(f" result = process_data({i})") lines.append(f" return result") return "\n".join(lines) def _long_text_content(): """生成超过 min_length 的纯文本内容""" return "This is plain text content. " * 100 # --------------------------------------------------------------------------- # TestHeadroomAvailability # --------------------------------------------------------------------------- class TestHeadroomAvailability: """测试 headroom-ai 可用性检测""" def test_is_available_false_when_not_installed(self): """_HEADROOM_AVAILABLE=False 时 is_available() 返回 False""" compressor = HeadroomCompressor({}) with patch("agentkit.core.headroom_compressor._HEADROOM_AVAILABLE", False): assert compressor.is_available() is False def test_is_available_true_when_installed(self): """_HEADROOM_AVAILABLE=True 时 is_available() 返回 True""" compressor = HeadroomCompressor({}) with patch("agentkit.core.headroom_compressor._HEADROOM_AVAILABLE", True): assert compressor.is_available() is True # --------------------------------------------------------------------------- # TestContentTypeDetection # --------------------------------------------------------------------------- class TestContentTypeDetection: """测试内容类型检测函数""" def test_json_content_detected(self): """有效 JSON 对象被正确检测""" assert _is_json_content('{"key": "value"}') is True def test_json_array_detected(self): """有效 JSON 数组被正确检测""" assert _is_json_content('[1, 2, 3]') is True def test_non_json_content(self): """普通文本不被识别为 JSON""" assert _is_json_content("hello world") is False def test_invalid_json_start(self): """以 { 开头但无效的 JSON 不被识别""" assert _is_json_content("{invalid") is False def test_code_content_detected(self): """Python 代码(含 def/class 关键字)被正确检测""" code = "def hello():\n pass\n\nclass Foo:\n pass\nimport os\nfrom sys import path" assert _is_code_content(code) is True def test_non_code_content(self): """纯文本不被识别为代码""" text = "This is just a regular paragraph of text with no code keywords at all." assert _is_code_content(text) is False # --------------------------------------------------------------------------- # TestCompressToolResult # --------------------------------------------------------------------------- class TestCompressToolResult: """测试 compress_tool_result 方法""" @pytest.mark.asyncio async def test_short_content_not_compressed(self): """短于 min_length 的内容不压缩""" compressor = HeadroomCompressor({"min_length": 500}) with patch("agentkit.core.headroom_compressor._HEADROOM_AVAILABLE", True): result = await compressor.compress_tool_result("test_tool", "short content") assert result == "short content" @pytest.mark.asyncio async def test_json_content_compressed_with_smart_crusher(self): """JSON 内容使用 smart_crusher 压缩""" compressor = HeadroomCompressor({ "min_length": 100, "compressors": ["smart_crusher", "code_compressor"], }) json_content = _long_json_content() mock_fn = MagicMock(return_value=_make_headroom_compress_mock("compressed json")) with patch("agentkit.core.headroom_compressor._HEADROOM_AVAILABLE", True), \ patch("agentkit.core.headroom_compressor.headroom_compress", mock_fn): result = await compressor.compress_tool_result("json_tool", json_content) assert "compressed json" in result assert "