|
|
"""
|
|
|
LLM生成器单元测试
|
|
|
|
|
|
本模块针对LLMGenerator类进行单元测试，验证API集成、重试机制、
|
|
|
流式功能和错误处理等核心功能。
|
|
|
"""
|
|
|
|
|
|
import os
|
|
|
import pytest
|
|
|
import asyncio
|
|
|
import aiohttp
|
|
|
import json
|
|
|
import time
|
|
|
from unittest.mock import Mock, patch, AsyncMock
|
|
|
from typing import Dict, Any, List
|
|
|
|
|
|
from src.spec.llm_generator import (
|
|
|
LLMGenerator, GenerationRequest, GenerationResult,
|
|
|
APIError, RateLimitError, TimeoutError, RetryExhaustedError,
|
|
|
generate_specification_sync, generate_batch_specifications_sync
|
|
|
)
|
|
|
from src.utils.logger import get_logger
|
|
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
|
|
|
class TestLLMGenerator:
|
|
|
"""LLM生成器单元测试类"""
|
|
|
|
|
|
@pytest.fixture
|
|
|
def mock_api_response(self):
|
|
|
"""模拟API响应"""
|
|
|
return {
|
|
|
"choices": [
|
|
|
{
|
|
|
"message": {
|
|
|
"content": "\\requires a >= 0 && b >= 0;\\n\\ensures return == a + b;"
|
|
|
}
|
|
|
}
|
|
|
],
|
|
|
"usage": {
|
|
|
"total_tokens": 150
|
|
|
}
|
|
|
}
|
|
|
|
|
|
@pytest.fixture
|
|
|
def mock_stream_response_chunks(self):
|
|
|
"""模拟流式响应块"""
|
|
|
return [
|
|
|
b'data: {"choices": [{"delta": {"content": "\\\\requires"}}]}\\n',
|
|
|
b'data: {"choices": [{"delta": {"content": " a >= 0"}}]}\\n',
|
|
|
b'data: {"choices": [{"delta": {"content": " && b >= 0;"}}]}\\n',
|
|
|
b'data: {"choices": [{"delta": {"content": "\\\\ensures"}}]}\\n',
|
|
|
b'data: {"choices": [{"delta": {"content": " return == a + b;"}}]}\\n',
|
|
|
b'data: [DONE]\\n'
|
|
|
]
|
|
|
|
|
|
@pytest.fixture
|
|
|
def sample_function_info(self):
|
|
|
"""示例函数信息"""
|
|
|
return {
|
|
|
'name': 'add',
|
|
|
'return_type': 'int',
|
|
|
'parameters': [
|
|
|
{'name': 'a', 'type': 'int'},
|
|
|
{'name': 'b', 'type': 'int'}
|
|
|
],
|
|
|
'variables': [],
|
|
|
'function_calls': []
|
|
|
}
|
|
|
|
|
|
@pytest.fixture
|
|
|
def sample_request(self, sample_function_info):
|
|
|
"""示例生成请求"""
|
|
|
return GenerationRequest(
|
|
|
function_name='add',
|
|
|
function_info=sample_function_info,
|
|
|
verification_goals=['functional_correctness', 'memory_safety'],
|
|
|
max_retries=2,
|
|
|
validate=True,
|
|
|
store=False
|
|
|
)
|
|
|
|
|
|
def test_llm_generator_initialization(self):
|
|
|
"""测试LLM生成器初始化"""
|
|
|
# 使用环境变量进行测试
|
|
|
with patch.dict(os.environ, {'SILICONFLOW_API_KEY': 'test_key'}):
|
|
|
generator = LLMGenerator()
|
|
|
assert generator.api_key == 'test_key'
|
|
|
assert generator.base_url == 'https://api.siliconflow.cn/v1'
|
|
|
assert generator.model == 'deepseek-ai/DeepSeek-V3.1'
|
|
|
assert generator.timeout > 0
|
|
|
assert generator.max_retries > 0
|
|
|
assert generator.temperature >= 0.0
|
|
|
assert generator.max_tokens > 0
|
|
|
|
|
|
def test_llm_generator_missing_api_key(self):
|
|
|
"""测试缺少API密钥时的错误处理"""
|
|
|
with patch.dict(os.environ, {}, clear=True):
|
|
|
with pytest.raises(Exception): # 应该抛出LLMError
|
|
|
LLMGenerator()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
|
async def test_context_manager(self):
|
|
|
"""测试异步上下文管理器"""
|
|
|
with patch.dict(os.environ, {'SILICONFLOW_API_KEY': 'test_key'}):
|
|
|
async with LLMGenerator() as generator:
|
|
|
assert generator.session is not None
|
|
|
assert generator._session_managed is False
|
|
|
|
|
|
# 退出上下文后session应该被关闭
|
|
|
assert generator.session is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
|
async def test_generate_specification_success(self, sample_request, mock_api_response):
|
|
|
"""测试成功的规范生成"""
|
|
|
with patch.dict(os.environ, {'SILICONFLOW_API_KEY': 'test_key'}):
|
|
|
with patch('aiohttp.ClientSession.post') as mock_post:
|
|
|
mock_response = AsyncMock()
|
|
|
mock_response.status = 200
|
|
|
mock_response.json = AsyncMock(return_value=mock_api_response)
|
|
|
mock_post.return_value.__aenter__.return_value = mock_response
|
|
|
|
|
|
async with LLMGenerator() as generator:
|
|
|
result = await generator.generate_specification(sample_request)
|
|
|
|
|
|
assert result.specification is not None
|
|
|
assert result.specification != ""
|
|
|
assert result.tokens_used == 150
|
|
|
assert result.generation_time > 0
|
|
|
assert result.quality_score >= 0.0
|
|
|
assert result.metadata['function_name'] == 'add'
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
|
async def test_generate_specification_api_error(self, sample_request):
|
|
|
"""测试API错误处理"""
|
|
|
with patch.dict(os.environ, {'SILICONFLOW_API_KEY': 'test_key'}):
|
|
|
with patch('aiohttp.ClientSession.post') as mock_post:
|
|
|
mock_response = AsyncMock()
|
|
|
mock_response.status = 500
|
|
|
mock_response.text = AsyncMock(return_value="Internal Server Error")
|
|
|
mock_post.return_value.__aenter__.return_value = mock_response
|
|
|
|
|
|
async with LLMGenerator() as generator:
|
|
|
with pytest.raises(APIError):
|
|
|
await generator.generate_specification(sample_request)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
|
async def test_generate_specification_rate_limit(self, sample_request):
|
|
|
"""测试速率限制处理"""
|
|
|
with patch.dict(os.environ, {'SILICONFLOW_API_KEY': 'test_key'}):
|
|
|
with patch('aiohttp.ClientSession.post') as mock_post:
|
|
|
mock_response = AsyncMock()
|
|
|
mock_response.status = 429
|
|
|
mock_post.return_value.__aenter__.return_value = mock_response
|
|
|
|
|
|
async with LLMGenerator() as generator:
|
|
|
with pytest.raises(RateLimitError):
|
|
|
await generator.generate_specification(sample_request)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
|
async def test_retry_mechanism(self, sample_request, mock_api_response):
|
|
|
"""测试重试机制"""
|
|
|
with patch.dict(os.environ, {'SILICONFLOW_API_KEY': 'test_key'}):
|
|
|
call_count = 0
|
|
|
|
|
|
async def mock_post(*args, **kwargs):
|
|
|
nonlocal call_count
|
|
|
call_count += 1
|
|
|
if call_count < 3: # 前两次失败
|
|
|
mock_response = AsyncMock()
|
|
|
mock_response.status = 500
|
|
|
mock_post.return_value.__aenter__.return_value = mock_response
|
|
|
return mock_response
|
|
|
else: # 第三次成功
|
|
|
mock_response = AsyncMock()
|
|
|
mock_response.status = 200
|
|
|
mock_response.json = AsyncMock(return_value=mock_api_response)
|
|
|
mock_post.return_value.__aenter__.return_value = mock_response
|
|
|
return mock_response
|
|
|
|
|
|
with patch('aiohttp.ClientSession.post', side_effect=mock_post):
|
|
|
async with LLMGenerator() as generator:
|
|
|
result = await generator.generate_specification(sample_request)
|
|
|
|
|
|
assert result.specification is not None
|
|
|
assert call_count == 3 # 应该重试了3次
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
|
async def test_retry_exhausted(self, sample_request):
|
|
|
"""测试重试次数耗尽"""
|
|
|
with patch.dict(os.environ, {'SILICONFLOW_API_KEY': 'test_key'}):
|
|
|
with patch('aiohttp.ClientSession.post') as mock_post:
|
|
|
mock_response = AsyncMock()
|
|
|
mock_response.status = 500
|
|
|
mock_post.return_value.__aenter__.return_value = mock_response
|
|
|
|
|
|
async with LLMGenerator() as generator:
|
|
|
request = GenerationRequest(
|
|
|
function_name='test',
|
|
|
function_info={'name': 'test'},
|
|
|
verification_goals=['test'],
|
|
|
max_retries=1 # 只重试1次
|
|
|
)
|
|
|
|
|
|
with pytest.raises(RetryExhaustedError):
|
|
|
await generator.generate_specification(request)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
|
async def test_stream_specification(self, sample_request, mock_stream_response_chunks):
|
|
|
"""测试流式规范生成"""
|
|
|
with patch.dict(os.environ, {'SILICONFLOW_API_KEY': 'test_key'}):
|
|
|
with patch('aiohttp.ClientSession.post') as mock_post:
|
|
|
mock_response = AsyncMock()
|
|
|
mock_response.status = 200
|
|
|
mock_response.content.iter_any = AsyncMock(return_value=mock_stream_response_chunks)
|
|
|
mock_post.return_value.__aenter__.return_value = mock_response
|
|
|
|
|
|
async with LLMGenerator() as generator:
|
|
|
chunks = []
|
|
|
async for chunk in generator.stream_specification(sample_request):
|
|
|
chunks.append(chunk)
|
|
|
|
|
|
assert len(chunks) > 0
|
|
|
full_spec = ''.join(chunks)
|
|
|
assert '\\requires' in full_spec
|
|
|
assert '\\ensures' in full_spec
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
|
async def test_batch_specifications(self, sample_request, mock_api_response):
|
|
|
"""测试批量规范生成"""
|
|
|
with patch.dict(os.environ, {'SILICONFLOW_API_KEY': 'test_key'}):
|
|
|
with patch('aiohttp.ClientSession.post') as mock_post:
|
|
|
mock_response = AsyncMock()
|
|
|
mock_response.status = 200
|
|
|
mock_response.json = AsyncMock(return_value=mock_api_response)
|
|
|
mock_post.return_value.__aenter__.return_value = mock_response
|
|
|
|
|
|
async with LLMGenerator() as generator:
|
|
|
requests = [sample_request, sample_request] # 两个相同的请求
|
|
|
results = await generator.generate_batch_specifications(requests)
|
|
|
|
|
|
assert len(results) == 2
|
|
|
for result in results:
|
|
|
assert result.specification is not None
|
|
|
assert result.tokens_used > 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
|
async def test_refine_specification(self, sample_request, mock_api_response):
|
|
|
"""测试规范细化功能"""
|
|
|
with patch.dict(os.environ, {'SILICONFLOW_API_KEY': 'test_key'}):
|
|
|
with patch('aiohttp.ClientSession.post') as mock_post:
|
|
|
mock_response = AsyncMock()
|
|
|
mock_response.status = 200
|
|
|
mock_response.json = AsyncMock(return_value=mock_api_response)
|
|
|
mock_post.return_value.__aenter__.return_value = mock_response
|
|
|
|
|
|
async with LLMGenerator() as generator:
|
|
|
original_spec = "\\requires a >= 0;\\n\\ensures return == a + b;"
|
|
|
issues = ["Missing parameter validation", "Insufficient coverage"]
|
|
|
suggestions = ["Add null checks", "Include overflow checks"]
|
|
|
|
|
|
result = await generator.refine_specification(
|
|
|
original_spec, sample_request.function_info, issues, suggestions
|
|
|
)
|
|
|
|
|
|
assert result.specification is not None
|
|
|
assert result.tokens_used > 0
|
|
|
assert result.quality_score >= 0.0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
|
async def test_health_check(self):
|
|
|
"""测试健康检查功能"""
|
|
|
with patch.dict(os.environ, {'SILICONFLOW_API_KEY': 'test_key'}):
|
|
|
with patch('aiohttp.ClientSession.post') as mock_post:
|
|
|
mock_response = AsyncMock()
|
|
|
mock_response.status = 200
|
|
|
mock_post.return_value.__aenter__.return_value = mock_response
|
|
|
|
|
|
async with LLMGenerator() as generator:
|
|
|
health = await generator.health_check()
|
|
|
|
|
|
assert health['status'] == 'healthy'
|
|
|
assert 'api_response_time' in health
|
|
|
assert 'timestamp' in health
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
|
async def test_health_check_failure(self):
|
|
|
"""测试健康检查失败情况"""
|
|
|
with patch.dict(os.environ, {'SILICONFLOW_API_KEY': 'test_key'}):
|
|
|
with patch('aiohttp.ClientSession.post') as mock_post:
|
|
|
mock_response = AsyncMock()
|
|
|
mock_response.status = 500
|
|
|
mock_post.return_value.__aenter__.return_value = mock_response
|
|
|
|
|
|
async with LLMGenerator() as generator:
|
|
|
health = await generator.health_check()
|
|
|
|
|
|
assert health['status'] == 'unhealthy'
|
|
|
|
|
|
def test_post_process_spec(self):
|
|
|
"""测试规范后处理功能"""
|
|
|
with patch.dict(os.environ, {'SILICONFLOW_API_KEY': 'test_key'}):
|
|
|
generator = LLMGenerator()
|
|
|
|
|
|
# 测试清理Markdown格式
|
|
|
markdown_spec = """
|
|
|
```c
|
|
|
\\requires a >= 0;
|
|
|
\\ensures return == a + b;
|
|
|
```
|
|
|
"""
|
|
|
cleaned = generator._post_process_spec(markdown_spec)
|
|
|
assert '\\requires' in cleaned
|
|
|
assert '\\ensures' in cleaned
|
|
|
assert '```' not in cleaned
|
|
|
|
|
|
# 测试空输入
|
|
|
empty_result = generator._post_process_spec("")
|
|
|
assert empty_result == ""
|
|
|
|
|
|
# 测试None输入
|
|
|
none_result = generator._post_process_spec(None)
|
|
|
assert none_result == ""
|
|
|
|
|
|
def test_generation_stats(self):
|
|
|
"""测试生成统计功能"""
|
|
|
with patch.dict(os.environ, {'SILICONFLOW_API_KEY': 'test_key'}):
|
|
|
generator = LLMGenerator()
|
|
|
|
|
|
# 模拟存储一些规范
|
|
|
with patch.object(generator.storage, 'list_specifications', return_value=[
|
|
|
{
|
|
|
'metadata': {
|
|
|
'quality_score': 0.8,
|
|
|
'model': 'test-model'
|
|
|
}
|
|
|
},
|
|
|
{
|
|
|
'metadata': {
|
|
|
'quality_score': 0.9,
|
|
|
'model': 'test-model'
|
|
|
}
|
|
|
}
|
|
|
]):
|
|
|
stats = generator.get_generation_stats()
|
|
|
|
|
|
assert stats['total_specifications'] == 2
|
|
|
assert stats['average_quality'] == 0.85
|
|
|
assert 'quality_distribution' in stats
|
|
|
assert 'model_usage' in stats
|
|
|
|
|
|
def test_synchronous_wrapper(self, sample_request):
|
|
|
"""测试同步包装器函数"""
|
|
|
with patch.dict(os.environ, {'SILICONFLOW_API_KEY': 'test_key'}):
|
|
|
with patch('aiohttp.ClientSession.post') as mock_post:
|
|
|
mock_response = AsyncMock()
|
|
|
mock_response.status = 200
|
|
|
mock_response.json = AsyncMock(return_value={
|
|
|
"choices": [{"message": {"content": "test spec"}}],
|
|
|
"usage": {"total_tokens": 100}
|
|
|
})
|
|
|
mock_post.return_value.__aenter__.return_value = mock_response
|
|
|
|
|
|
# 测试同步函数
|
|
|
result = generate_specification_sync(sample_request)
|
|
|
assert result.specification == "test spec"
|
|
|
assert result.tokens_used == 100
|
|
|
|
|
|
def test_batch_synchronous_wrapper(self, sample_request):
|
|
|
"""测试批量同步包装器函数"""
|
|
|
with patch.dict(os.environ, {'SILICONFLOW_API_KEY': 'test_key'}):
|
|
|
with patch('aiohttp.ClientSession.post') as mock_post:
|
|
|
mock_response = AsyncMock()
|
|
|
mock_response.status = 200
|
|
|
mock_response.json = AsyncMock(return_value={
|
|
|
"choices": [{"message": {"content": "test spec"}}],
|
|
|
"usage": {"total_tokens": 100}
|
|
|
})
|
|
|
mock_post.return_value.__aenter__.return_value = mock_response
|
|
|
|
|
|
# 测试批量同步函数
|
|
|
requests = [sample_request, sample_request]
|
|
|
results = generate_batch_specifications_sync(requests)
|
|
|
|
|
|
assert len(results) == 2
|
|
|
for result in results:
|
|
|
assert result.specification == "test spec"
|
|
|
assert result.tokens_used == 100
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
|
async def test_timeout_handling(self, sample_request):
|
|
|
"""测试超时处理"""
|
|
|
with patch.dict(os.environ, {'SILICONFLOW_API_KEY': 'test_key'}):
|
|
|
with patch('aiohttp.ClientSession.post') as mock_post:
|
|
|
# 模拟超时
|
|
|
mock_post.side_effect = asyncio.TimeoutError()
|
|
|
|
|
|
async with LLMGenerator() as generator:
|
|
|
with pytest.raises(TimeoutError):
|
|
|
await generator.generate_specification(sample_request)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
|
async def test_network_error_handling(self, sample_request):
|
|
|
"""测试网络错误处理"""
|
|
|
with patch.dict(os.environ, {'SILICONFLOW_API_KEY': 'test_key'}):
|
|
|
with patch('aiohttp.ClientSession.post') as mock_post:
|
|
|
# 模拟网络错误
|
|
|
mock_post.side_effect = aiohttp.ClientError("Network error")
|
|
|
|
|
|
async with LLMGenerator() as generator:
|
|
|
with pytest.raises(APIError):
|
|
|
await generator.generate_specification(sample_request)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
|
async def test_invalid_json_response(self, sample_request):
|
|
|
"""测试无效JSON响应处理"""
|
|
|
with patch.dict(os.environ, {'SILICONFLOW_API_KEY': 'test_key'}):
|
|
|
with patch('aiohttp.ClientSession.post') as mock_post:
|
|
|
mock_response = AsyncMock()
|
|
|
mock_response.status = 200
|
|
|
mock_response.json = AsyncMock(side_effect=json.JSONDecodeError("Invalid JSON", "", 0))
|
|
|
mock_post.return_value.__aenter__.return_value = mock_response
|
|
|
|
|
|
async with LLMGenerator() as generator:
|
|
|
with pytest.raises(APIError):
|
|
|
await generator.generate_specification(sample_request)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
|
async def test_missing_choices_in_response(self, sample_request):
|
|
|
"""测试响应中缺少choices字段"""
|
|
|
with patch.dict(os.environ, {'SILICONFLOW_API_KEY': 'test_key'}):
|
|
|
with patch('aiohttp.ClientSession.post') as mock_post:
|
|
|
mock_response = AsyncMock()
|
|
|
mock_response.status = 200
|
|
|
mock_response.json = AsyncMock(return_value={}) # 空响应
|
|
|
mock_post.return_value.__aenter__.return_value = mock_response
|
|
|
|
|
|
async with LLMGenerator() as generator:
|
|
|
with pytest.raises(APIError):
|
|
|
await generator.generate_specification(sample_request)
|
|
|
|
|
|
def test_session_management(self):
|
|
|
"""测试会话管理"""
|
|
|
with patch.dict(os.environ, {'SILICONFLOW_API_KEY': 'test_key'}):
|
|
|
generator = LLMGenerator()
|
|
|
|
|
|
# 测试自动创建会话
|
|
|
assert generator.session is None
|
|
|
generator._ensure_session()
|
|
|
assert generator.session is not None
|
|
|
assert generator._session_managed is True
|
|
|
|
|
|
# 测试清理会话
|
|
|
loop = asyncio.new_event_loop()
|
|
|
asyncio.set_event_loop(loop)
|
|
|
loop.run_until_complete(generator._cleanup_session())
|
|
|
assert generator.session is None
|
|
|
loop.close()
|
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
pytest.main([__file__, "-v"]) |
