427 lines
13 KiB
Python
427 lines
13 KiB
Python
"""单元测试：错误处理机制。"""
|
||
|
||
import pytest
|
||
import pandas as pd
|
||
import time
|
||
from pathlib import Path
|
||
from unittest.mock import Mock, patch, MagicMock
|
||
import tempfile
|
||
import os
|
||
|
||
from src.error_handling import (
|
||
load_data_with_retry,
|
||
call_llm_with_fallback,
|
||
execute_tool_safely,
|
||
execute_task_with_recovery,
|
||
validate_tool_params,
|
||
validate_tool_result,
|
||
DataLoadError,
|
||
AICallError,
|
||
ToolExecutionError
|
||
)
|
||
|
||
|
||
class TestLoadDataWithRetry:
|
||
"""测试数据加载错误处理。"""
|
||
|
||
def test_load_valid_csv(self, tmp_path):
|
||
"""测试加载有效的 CSV 文件。"""
|
||
# 创建测试文件
|
||
csv_file = tmp_path / "test.csv"
|
||
df = pd.DataFrame({
|
||
'col1': [1, 2, 3],
|
||
'col2': ['a', 'b', 'c']
|
||
})
|
||
df.to_csv(csv_file, index=False)
|
||
|
||
# 加载数据
|
||
result = load_data_with_retry(str(csv_file))
|
||
|
||
assert len(result) == 3
|
||
assert len(result.columns) == 2
|
||
assert list(result.columns) == ['col1', 'col2']
|
||
|
||
def test_load_gbk_encoded_file(self, tmp_path):
|
||
"""测试加载 GBK 编码的文件。"""
|
||
# 创建 GBK 编码的文件
|
||
csv_file = tmp_path / "test_gbk.csv"
|
||
df = pd.DataFrame({
|
||
'列1': [1, 2, 3],
|
||
'列2': ['中文', '测试', '数据']
|
||
})
|
||
df.to_csv(csv_file, index=False, encoding='gbk')
|
||
|
||
# 加载数据
|
||
result = load_data_with_retry(str(csv_file))
|
||
|
||
assert len(result) == 3
|
||
assert '列1' in result.columns
|
||
assert '列2' in result.columns
|
||
|
||
def test_load_file_not_exists(self):
|
||
"""测试文件不存在的情况。"""
|
||
with pytest.raises(DataLoadError, match="文件不存在"):
|
||
load_data_with_retry("nonexistent.csv")
|
||
|
||
def test_load_empty_file(self, tmp_path):
|
||
"""测试空文件的处理。"""
|
||
# 创建空文件
|
||
csv_file = tmp_path / "empty.csv"
|
||
csv_file.touch()
|
||
|
||
with pytest.raises(DataLoadError, match="文件为空"):
|
||
load_data_with_retry(str(csv_file))
|
||
|
||
def test_load_large_file_sampling(self, tmp_path):
|
||
"""测试大文件采样。"""
|
||
# 创建大文件（模拟）
|
||
csv_file = tmp_path / "large.csv"
|
||
df = pd.DataFrame({
|
||
'col1': range(2000000),
|
||
'col2': range(2000000)
|
||
})
|
||
# 只保存前 1500000 行以加快测试
|
||
df.head(1500000).to_csv(csv_file, index=False)
|
||
|
||
# 加载数据（应该采样到 1000000 行）
|
||
result = load_data_with_retry(str(csv_file), sample_size=1000000)
|
||
|
||
assert len(result) == 1000000
|
||
|
||
def test_load_different_separator(self, tmp_path):
|
||
"""测试不同分隔符的文件。"""
|
||
# 创建使用分号分隔的文件
|
||
csv_file = tmp_path / "semicolon.csv"
|
||
with open(csv_file, 'w') as f:
|
||
f.write("col1;col2\n")
|
||
f.write("1;a\n")
|
||
f.write("2;b\n")
|
||
|
||
# 加载数据
|
||
result = load_data_with_retry(str(csv_file))
|
||
|
||
assert len(result) == 2
|
||
assert len(result.columns) == 2
|
||
|
||
|
||
class TestCallLLMWithFallback:
|
||
"""测试 AI 调用错误处理。"""
|
||
|
||
def test_successful_call(self):
|
||
"""测试成功的 AI 调用。"""
|
||
mock_func = Mock(return_value={'result': 'success'})
|
||
|
||
result = call_llm_with_fallback(mock_func, prompt="test")
|
||
|
||
assert result == {'result': 'success'}
|
||
assert mock_func.call_count == 1
|
||
|
||
def test_retry_on_timeout(self):
|
||
"""测试超时重试机制。"""
|
||
mock_func = Mock(side_effect=[
|
||
TimeoutError("timeout"),
|
||
TimeoutError("timeout"),
|
||
{'result': 'success'}
|
||
])
|
||
|
||
result = call_llm_with_fallback(mock_func, max_retries=3, prompt="test")
|
||
|
||
assert result == {'result': 'success'}
|
||
assert mock_func.call_count == 3
|
||
|
||
def test_exponential_backoff(self):
|
||
"""测试指数退避。"""
|
||
mock_func = Mock(side_effect=[
|
||
Exception("error"),
|
||
{'result': 'success'}
|
||
])
|
||
|
||
start_time = time.time()
|
||
result = call_llm_with_fallback(mock_func, max_retries=3, prompt="test")
|
||
elapsed = time.time() - start_time
|
||
|
||
# 应该等待至少 1 秒（2^0）
|
||
assert elapsed >= 1.0
|
||
assert result == {'result': 'success'}
|
||
|
||
def test_fallback_on_failure(self):
|
||
"""测试降级策略。"""
|
||
mock_func = Mock(side_effect=Exception("error"))
|
||
fallback_func = Mock(return_value={'result': 'fallback'})
|
||
|
||
result = call_llm_with_fallback(
|
||
mock_func,
|
||
fallback_func=fallback_func,
|
||
max_retries=2,
|
||
prompt="test"
|
||
)
|
||
|
||
assert result == {'result': 'fallback'}
|
||
assert mock_func.call_count == 2
|
||
assert fallback_func.call_count == 1
|
||
|
||
def test_no_fallback_raises_error(self):
|
||
"""测试无降级策略时抛出错误。"""
|
||
mock_func = Mock(side_effect=Exception("error"))
|
||
|
||
with pytest.raises(AICallError, match="AI 调用失败"):
|
||
call_llm_with_fallback(mock_func, max_retries=2, prompt="test")
|
||
|
||
def test_fallback_also_fails(self):
|
||
"""测试降级策略也失败的情况。"""
|
||
mock_func = Mock(side_effect=Exception("error"))
|
||
fallback_func = Mock(side_effect=Exception("fallback error"))
|
||
|
||
with pytest.raises(AICallError, match="AI 调用和降级策略都失败"):
|
||
call_llm_with_fallback(
|
||
mock_func,
|
||
fallback_func=fallback_func,
|
||
max_retries=2,
|
||
prompt="test"
|
||
)
|
||
|
||
|
||
class TestExecuteToolSafely:
|
||
"""测试工具执行错误处理。"""
|
||
|
||
def test_successful_execution(self):
|
||
"""测试成功的工具执行。"""
|
||
mock_tool = Mock()
|
||
mock_tool.name = "test_tool"
|
||
mock_tool.parameters = {'required': [], 'properties': {}}
|
||
mock_tool.execute = Mock(return_value={'data': 'result'})
|
||
|
||
df = pd.DataFrame({'col1': [1, 2, 3]})
|
||
result = execute_tool_safely(mock_tool, df)
|
||
|
||
assert result['success'] is True
|
||
assert result['data'] == {'data': 'result'}
|
||
assert result['tool'] == 'test_tool'
|
||
|
||
def test_missing_execute_method(self):
|
||
"""测试工具缺少 execute 方法。"""
|
||
mock_tool = Mock(spec=[])
|
||
mock_tool.name = "bad_tool"
|
||
|
||
df = pd.DataFrame({'col1': [1, 2, 3]})
|
||
result = execute_tool_safely(mock_tool, df)
|
||
|
||
assert result['success'] is False
|
||
assert 'execute 方法' in result['error']
|
||
|
||
def test_parameter_validation_failure(self):
|
||
"""测试参数验证失败。"""
|
||
mock_tool = Mock()
|
||
mock_tool.name = "test_tool"
|
||
mock_tool.parameters = {
|
||
'required': ['column'],
|
||
'properties': {
|
||
'column': {'type': 'string'}
|
||
}
|
||
}
|
||
mock_tool.execute = Mock(return_value={'data': 'result'})
|
||
|
||
df = pd.DataFrame({'col1': [1, 2, 3]})
|
||
# 缺少必需参数
|
||
result = execute_tool_safely(mock_tool, df)
|
||
|
||
assert result['success'] is False
|
||
assert '参数验证失败' in result['error']
|
||
|
||
def test_empty_data(self):
|
||
"""测试空数据。"""
|
||
mock_tool = Mock()
|
||
mock_tool.name = "test_tool"
|
||
mock_tool.parameters = {'required': [], 'properties': {}}
|
||
|
||
df = pd.DataFrame()
|
||
result = execute_tool_safely(mock_tool, df)
|
||
|
||
assert result['success'] is False
|
||
assert '数据为空' in result['error']
|
||
|
||
def test_execution_exception(self):
|
||
"""测试执行异常。"""
|
||
mock_tool = Mock()
|
||
mock_tool.name = "test_tool"
|
||
mock_tool.parameters = {'required': [], 'properties': {}}
|
||
mock_tool.execute = Mock(side_effect=Exception("execution error"))
|
||
|
||
df = pd.DataFrame({'col1': [1, 2, 3]})
|
||
result = execute_tool_safely(mock_tool, df)
|
||
|
||
assert result['success'] is False
|
||
assert 'execution error' in result['error']
|
||
|
||
|
||
class TestValidateToolParams:
|
||
"""测试工具参数验证。"""
|
||
|
||
def test_valid_params(self):
|
||
"""测试有效参数。"""
|
||
mock_tool = Mock()
|
||
mock_tool.parameters = {
|
||
'required': ['column'],
|
||
'properties': {
|
||
'column': {'type': 'string'}
|
||
}
|
||
}
|
||
|
||
result = validate_tool_params(mock_tool, {'column': 'col1'})
|
||
|
||
assert result['valid'] is True
|
||
|
||
def test_missing_required_param(self):
|
||
"""测试缺少必需参数。"""
|
||
mock_tool = Mock()
|
||
mock_tool.parameters = {
|
||
'required': ['column'],
|
||
'properties': {}
|
||
}
|
||
|
||
result = validate_tool_params(mock_tool, {})
|
||
|
||
assert result['valid'] is False
|
||
assert '缺少必需参数' in result['error']
|
||
|
||
def test_wrong_param_type(self):
|
||
"""测试参数类型错误。"""
|
||
mock_tool = Mock()
|
||
mock_tool.parameters = {
|
||
'required': [],
|
||
'properties': {
|
||
'column': {'type': 'string'}
|
||
}
|
||
}
|
||
|
||
result = validate_tool_params(mock_tool, {'column': 123})
|
||
|
||
assert result['valid'] is False
|
||
assert '应为字符串类型' in result['error']
|
||
|
||
|
||
class TestValidateToolResult:
|
||
"""测试工具结果验证。"""
|
||
|
||
def test_valid_result(self):
|
||
"""测试有效结果。"""
|
||
result = validate_tool_result({'data': 'test'})
|
||
|
||
assert result['valid'] is True
|
||
|
||
def test_none_result(self):
|
||
"""测试 None 结果。"""
|
||
result = validate_tool_result(None)
|
||
|
||
assert result['valid'] is False
|
||
assert 'None' in result['error']
|
||
|
||
def test_wrong_type_result(self):
|
||
"""测试错误类型结果。"""
|
||
result = validate_tool_result("string result")
|
||
|
||
assert result['valid'] is False
|
||
assert '类型错误' in result['error']
|
||
|
||
|
||
class TestExecuteTaskWithRecovery:
|
||
"""测试任务执行错误处理。"""
|
||
|
||
def test_successful_execution(self):
|
||
"""测试成功的任务执行。"""
|
||
mock_task = Mock()
|
||
mock_task.id = "task1"
|
||
mock_task.name = "Test Task"
|
||
mock_task.dependencies = []
|
||
|
||
mock_plan = Mock()
|
||
mock_plan.tasks = [mock_task]
|
||
|
||
mock_execute = Mock(return_value=Mock(success=True))
|
||
|
||
result = execute_task_with_recovery(mock_task, mock_plan, mock_execute)
|
||
|
||
assert mock_task.status == 'completed'
|
||
assert mock_execute.call_count == 1
|
||
|
||
def test_skip_on_missing_dependency(self):
|
||
"""测试依赖任务不存在时跳过。"""
|
||
mock_task = Mock()
|
||
mock_task.id = "task2"
|
||
mock_task.name = "Test Task"
|
||
mock_task.dependencies = ["task1"]
|
||
|
||
mock_plan = Mock()
|
||
mock_plan.tasks = [mock_task]
|
||
|
||
mock_execute = Mock()
|
||
|
||
result = execute_task_with_recovery(mock_task, mock_plan, mock_execute)
|
||
|
||
assert mock_task.status == 'skipped'
|
||
assert mock_execute.call_count == 0
|
||
|
||
def test_skip_on_failed_dependency(self):
|
||
"""测试依赖任务失败时跳过。"""
|
||
mock_dep_task = Mock()
|
||
mock_dep_task.id = "task1"
|
||
mock_dep_task.status = 'failed'
|
||
|
||
mock_task = Mock()
|
||
mock_task.id = "task2"
|
||
mock_task.name = "Test Task"
|
||
mock_task.dependencies = ["task1"]
|
||
|
||
mock_plan = Mock()
|
||
mock_plan.tasks = [mock_dep_task, mock_task]
|
||
|
||
mock_execute = Mock()
|
||
|
||
result = execute_task_with_recovery(mock_task, mock_plan, mock_execute)
|
||
|
||
assert mock_task.status == 'skipped'
|
||
assert mock_execute.call_count == 0
|
||
|
||
def test_mark_failed_on_exception(self):
|
||
"""测试执行异常时标记失败。"""
|
||
mock_task = Mock()
|
||
mock_task.id = "task1"
|
||
mock_task.name = "Test Task"
|
||
mock_task.dependencies = []
|
||
|
||
mock_plan = Mock()
|
||
mock_plan.tasks = [mock_task]
|
||
|
||
mock_execute = Mock(side_effect=Exception("execution error"))
|
||
|
||
result = execute_task_with_recovery(mock_task, mock_plan, mock_execute)
|
||
|
||
assert mock_task.status == 'failed'
|
||
|
||
def test_continue_on_task_failure(self):
|
||
"""测试单个任务失败不影响其他任务。"""
|
||
mock_task1 = Mock()
|
||
mock_task1.id = "task1"
|
||
mock_task1.name = "Task 1"
|
||
mock_task1.dependencies = []
|
||
|
||
mock_task2 = Mock()
|
||
mock_task2.id = "task2"
|
||
mock_task2.name = "Task 2"
|
||
mock_task2.dependencies = []
|
||
|
||
mock_plan = Mock()
|
||
mock_plan.tasks = [mock_task1, mock_task2]
|
||
|
||
# 第一个任务失败
|
||
mock_execute = Mock(side_effect=Exception("error"))
|
||
result1 = execute_task_with_recovery(mock_task1, mock_plan, mock_execute)
|
||
|
||
assert mock_task1.status == 'failed'
|
||
|
||
# 第二个任务应该可以继续执行
|
||
mock_execute2 = Mock(return_value=Mock(success=True))
|
||
result2 = execute_task_with_recovery(mock_task2, mock_plan, mock_execute2)
|
||
|
||
assert mock_task2.status == 'completed'
|