681 lines
22 KiB
Python
681 lines
22 KiB
Python
"""工具系统的单元测试。"""
|
||
|
||
import pytest
|
||
import pandas as pd
|
||
import numpy as np
|
||
from datetime import datetime, timedelta
|
||
|
||
from src.tools.base import AnalysisTool, ToolRegistry
|
||
from src.tools.query_tools import (
|
||
GetColumnDistributionTool,
|
||
GetValueCountsTool,
|
||
GetTimeSeriesTool,
|
||
GetCorrelationTool
|
||
)
|
||
from src.tools.stats_tools import (
|
||
CalculateStatisticsTool,
|
||
PerformGroupbyTool,
|
||
DetectOutliersTool,
|
||
CalculateTrendTool
|
||
)
|
||
from src.models import DataProfile, ColumnInfo
|
||
|
||
|
||
class TestGetColumnDistributionTool:
|
||
"""测试列分布工具。"""
|
||
|
||
def test_basic_functionality(self):
|
||
"""测试基本功能。"""
|
||
tool = GetColumnDistributionTool()
|
||
df = pd.DataFrame({
|
||
'status': ['open', 'closed', 'open', 'pending', 'closed', 'open']
|
||
})
|
||
|
||
result = tool.execute(df, column='status')
|
||
|
||
assert 'distribution' in result
|
||
assert result['column'] == 'status'
|
||
assert result['total_count'] == 6
|
||
assert result['unique_count'] == 3
|
||
assert len(result['distribution']) == 3
|
||
|
||
def test_top_n_limit(self):
|
||
"""测试 top_n 参数限制。"""
|
||
tool = GetColumnDistributionTool()
|
||
df = pd.DataFrame({
|
||
'value': list(range(20))
|
||
})
|
||
|
||
result = tool.execute(df, column='value', top_n=5)
|
||
|
||
assert len(result['distribution']) == 5
|
||
|
||
def test_nonexistent_column(self):
|
||
"""测试不存在的列。"""
|
||
tool = GetColumnDistributionTool()
|
||
df = pd.DataFrame({'col1': [1, 2, 3]})
|
||
|
||
result = tool.execute(df, column='nonexistent')
|
||
|
||
assert 'error' in result
|
||
|
||
|
||
class TestGetValueCountsTool:
|
||
"""测试值计数工具。"""
|
||
|
||
def test_basic_functionality(self):
|
||
"""测试基本功能。"""
|
||
tool = GetValueCountsTool()
|
||
df = pd.DataFrame({
|
||
'category': ['A', 'B', 'A', 'C', 'B', 'A']
|
||
})
|
||
|
||
result = tool.execute(df, column='category')
|
||
|
||
assert 'value_counts' in result
|
||
assert result['value_counts']['A'] == 3
|
||
assert result['value_counts']['B'] == 2
|
||
assert result['value_counts']['C'] == 1
|
||
|
||
def test_normalized_counts(self):
|
||
"""测试归一化计数。"""
|
||
tool = GetValueCountsTool()
|
||
df = pd.DataFrame({
|
||
'category': ['A', 'A', 'B', 'B']
|
||
})
|
||
|
||
result = tool.execute(df, column='category', normalize=True)
|
||
|
||
assert result['normalized'] is True
|
||
assert abs(result['value_counts']['A'] - 0.5) < 0.01
|
||
assert abs(result['value_counts']['B'] - 0.5) < 0.01
|
||
|
||
|
||
class TestGetTimeSeriesTool:
|
||
"""测试时间序列工具。"""
|
||
|
||
def test_basic_functionality(self):
|
||
"""测试基本功能。"""
|
||
tool = GetTimeSeriesTool()
|
||
dates = pd.date_range('2020-01-01', periods=10, freq='D')
|
||
df = pd.DataFrame({
|
||
'date': dates,
|
||
'value': range(10)
|
||
})
|
||
|
||
result = tool.execute(df, time_column='date', value_column='value', aggregation='sum')
|
||
|
||
assert 'time_series' in result
|
||
assert result['time_column'] == 'date'
|
||
assert result['aggregation'] == 'sum'
|
||
assert len(result['time_series']) > 0
|
||
|
||
def test_count_aggregation(self):
|
||
"""测试计数聚合。"""
|
||
tool = GetTimeSeriesTool()
|
||
dates = pd.date_range('2020-01-01', periods=5, freq='D')
|
||
df = pd.DataFrame({'date': dates})
|
||
|
||
result = tool.execute(df, time_column='date', aggregation='count')
|
||
|
||
assert 'time_series' in result
|
||
assert len(result['time_series']) > 0
|
||
|
||
def test_output_limit(self):
|
||
"""测试输出限制（不超过100行）。"""
|
||
tool = GetTimeSeriesTool()
|
||
dates = pd.date_range('2020-01-01', periods=200, freq='D')
|
||
df = pd.DataFrame({'date': dates})
|
||
|
||
result = tool.execute(df, time_column='date')
|
||
|
||
assert len(result['time_series']) <= 100
|
||
assert result['total_points'] == 200
|
||
assert result['returned_points'] == 100
|
||
|
||
|
||
class TestGetCorrelationTool:
|
||
"""测试相关性分析工具。"""
|
||
|
||
def test_basic_functionality(self):
|
||
"""测试基本功能。"""
|
||
tool = GetCorrelationTool()
|
||
df = pd.DataFrame({
|
||
'x': [1, 2, 3, 4, 5],
|
||
'y': [2, 4, 6, 8, 10],
|
||
'z': [1, 1, 1, 1, 1]
|
||
})
|
||
|
||
result = tool.execute(df)
|
||
|
||
assert 'correlation_matrix' in result
|
||
assert 'x' in result['correlation_matrix']
|
||
assert 'y' in result['correlation_matrix']
|
||
# x 和 y 完全正相关
|
||
assert abs(result['correlation_matrix']['x']['y'] - 1.0) < 0.01
|
||
|
||
def test_insufficient_numeric_columns(self):
|
||
"""测试数值列不足的情况。"""
|
||
tool = GetCorrelationTool()
|
||
df = pd.DataFrame({
|
||
'x': [1, 2, 3],
|
||
'text': ['a', 'b', 'c']
|
||
})
|
||
|
||
result = tool.execute(df)
|
||
|
||
assert 'error' in result
|
||
|
||
|
||
class TestCalculateStatisticsTool:
|
||
"""测试统计计算工具。"""
|
||
|
||
def test_basic_functionality(self):
|
||
"""测试基本功能。"""
|
||
tool = CalculateStatisticsTool()
|
||
df = pd.DataFrame({
|
||
'values': [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
|
||
})
|
||
|
||
result = tool.execute(df, column='values')
|
||
|
||
assert result['mean'] == 5.5
|
||
assert result['median'] == 5.5
|
||
assert result['min'] == 1
|
||
assert result['max'] == 10
|
||
assert result['count'] == 10
|
||
|
||
def test_non_numeric_column(self):
|
||
"""测试非数值列。"""
|
||
tool = CalculateStatisticsTool()
|
||
df = pd.DataFrame({
|
||
'text': ['a', 'b', 'c']
|
||
})
|
||
|
||
result = tool.execute(df, column='text')
|
||
|
||
assert 'error' in result
|
||
|
||
|
||
class TestPerformGroupbyTool:
|
||
"""测试分组聚合工具。"""
|
||
|
||
def test_basic_functionality(self):
|
||
"""测试基本功能。"""
|
||
tool = PerformGroupbyTool()
|
||
df = pd.DataFrame({
|
||
'category': ['A', 'B', 'A', 'B', 'A'],
|
||
'value': [10, 20, 30, 40, 50]
|
||
})
|
||
|
||
result = tool.execute(df, group_by='category', value_column='value', aggregation='sum')
|
||
|
||
assert 'groups' in result
|
||
assert len(result['groups']) == 2
|
||
# 找到 A 组的总和
|
||
group_a = next(g for g in result['groups'] if g['group'] == 'A')
|
||
assert group_a['value'] == 90 # 10 + 30 + 50
|
||
|
||
def test_count_aggregation(self):
|
||
"""测试计数聚合。"""
|
||
tool = PerformGroupbyTool()
|
||
df = pd.DataFrame({
|
||
'category': ['A', 'B', 'A', 'B', 'A']
|
||
})
|
||
|
||
result = tool.execute(df, group_by='category')
|
||
|
||
assert len(result['groups']) == 2
|
||
group_a = next(g for g in result['groups'] if g['group'] == 'A')
|
||
assert group_a['value'] == 3
|
||
|
||
def test_output_limit(self):
|
||
"""测试输出限制（不超过100组）。"""
|
||
tool = PerformGroupbyTool()
|
||
df = pd.DataFrame({
|
||
'category': [f'cat_{i}' for i in range(200)],
|
||
'value': range(200)
|
||
})
|
||
|
||
result = tool.execute(df, group_by='category', value_column='value', aggregation='sum')
|
||
|
||
assert len(result['groups']) <= 100
|
||
assert result['total_groups'] == 200
|
||
assert result['returned_groups'] == 100
|
||
|
||
|
||
class TestDetectOutliersTool:
|
||
"""测试异常值检测工具。"""
|
||
|
||
def test_iqr_method(self):
|
||
"""测试 IQR 方法。"""
|
||
tool = DetectOutliersTool()
|
||
# 创建包含明显异常值的数据
|
||
df = pd.DataFrame({
|
||
'values': [1, 2, 3, 4, 5, 6, 7, 8, 9, 100]
|
||
})
|
||
|
||
result = tool.execute(df, column='values', method='iqr')
|
||
|
||
assert result['outlier_count'] > 0
|
||
assert 100 in result['outlier_values']
|
||
|
||
def test_zscore_method(self):
|
||
"""测试 Z-score 方法。"""
|
||
tool = DetectOutliersTool()
|
||
df = pd.DataFrame({
|
||
'values': [1, 2, 3, 4, 5, 6, 7, 8, 9, 100]
|
||
})
|
||
|
||
result = tool.execute(df, column='values', method='zscore', threshold=2)
|
||
|
||
assert result['outlier_count'] > 0
|
||
assert result['method'] == 'zscore'
|
||
|
||
|
||
class TestCalculateTrendTool:
|
||
"""测试趋势计算工具。"""
|
||
|
||
def test_increasing_trend(self):
|
||
"""测试上升趋势。"""
|
||
tool = CalculateTrendTool()
|
||
dates = pd.date_range('2020-01-01', periods=10, freq='D')
|
||
df = pd.DataFrame({
|
||
'date': dates,
|
||
'value': range(10)
|
||
})
|
||
|
||
result = tool.execute(df, time_column='date', value_column='value')
|
||
|
||
assert result['trend'] == 'increasing'
|
||
assert result['slope'] > 0
|
||
assert result['r_squared'] > 0.9 # 完美线性关系
|
||
|
||
def test_decreasing_trend(self):
|
||
"""测试下降趋势。"""
|
||
tool = CalculateTrendTool()
|
||
dates = pd.date_range('2020-01-01', periods=10, freq='D')
|
||
df = pd.DataFrame({
|
||
'date': dates,
|
||
'value': list(range(10, 0, -1))
|
||
})
|
||
|
||
result = tool.execute(df, time_column='date', value_column='value')
|
||
|
||
assert result['trend'] == 'decreasing'
|
||
assert result['slope'] < 0
|
||
|
||
|
||
class TestToolParameterValidation:
|
||
"""测试工具参数验证。"""
|
||
|
||
def test_missing_required_parameter(self):
|
||
"""测试缺少必需参数。"""
|
||
tool = GetColumnDistributionTool()
|
||
df = pd.DataFrame({'col': [1, 2, 3]})
|
||
|
||
# 不提供必需的 column 参数
|
||
result = tool.execute(df)
|
||
|
||
# 应该返回错误或引发异常
|
||
assert 'error' in result or result is None
|
||
|
||
def test_invalid_aggregation_method(self):
|
||
"""测试无效的聚合方法。"""
|
||
tool = PerformGroupbyTool()
|
||
df = pd.DataFrame({
|
||
'category': ['A', 'B'],
|
||
'value': [1, 2]
|
||
})
|
||
|
||
result = tool.execute(df, group_by='category', value_column='value', aggregation='invalid')
|
||
|
||
assert 'error' in result
|
||
|
||
|
||
class TestToolErrorHandling:
|
||
"""测试工具错误处理。"""
|
||
|
||
def test_empty_dataframe(self):
|
||
"""测试空 DataFrame。"""
|
||
tool = CalculateStatisticsTool()
|
||
df = pd.DataFrame()
|
||
|
||
result = tool.execute(df, column='nonexistent')
|
||
|
||
assert 'error' in result
|
||
|
||
def test_all_null_values(self):
|
||
"""测试全部为空值的列。"""
|
||
tool = CalculateStatisticsTool()
|
||
df = pd.DataFrame({
|
||
'values': [None, None, None]
|
||
})
|
||
|
||
result = tool.execute(df, column='values')
|
||
|
||
# 应该处理空值情况
|
||
assert 'error' in result or result['count'] == 0
|
||
|
||
def test_invalid_date_column(self):
|
||
"""测试无效的日期列。"""
|
||
tool = GetTimeSeriesTool()
|
||
df = pd.DataFrame({
|
||
'not_date': ['a', 'b', 'c']
|
||
})
|
||
|
||
result = tool.execute(df, time_column='not_date')
|
||
|
||
assert 'error' in result
|
||
|
||
|
||
class TestToolRegistry:
|
||
"""测试工具注册表。"""
|
||
|
||
def test_register_and_retrieve(self):
|
||
"""测试注册和检索工具。"""
|
||
registry = ToolRegistry()
|
||
tool = GetColumnDistributionTool()
|
||
|
||
registry.register(tool)
|
||
retrieved = registry.get_tool(tool.name)
|
||
|
||
assert retrieved.name == tool.name
|
||
|
||
def test_unregister(self):
|
||
"""测试注销工具。"""
|
||
registry = ToolRegistry()
|
||
tool = GetColumnDistributionTool()
|
||
|
||
registry.register(tool)
|
||
registry.unregister(tool.name)
|
||
|
||
with pytest.raises(KeyError):
|
||
registry.get_tool(tool.name)
|
||
|
||
def test_list_tools(self):
|
||
"""测试列出所有工具。"""
|
||
registry = ToolRegistry()
|
||
tool1 = GetColumnDistributionTool()
|
||
tool2 = GetValueCountsTool()
|
||
|
||
registry.register(tool1)
|
||
registry.register(tool2)
|
||
|
||
tools = registry.list_tools()
|
||
assert len(tools) == 2
|
||
assert tool1.name in tools
|
||
assert tool2.name in tools
|
||
|
||
def test_get_applicable_tools(self):
|
||
"""测试获取适用的工具。"""
|
||
registry = ToolRegistry()
|
||
|
||
# 注册所有工具
|
||
registry.register(GetColumnDistributionTool())
|
||
registry.register(CalculateStatisticsTool())
|
||
registry.register(GetTimeSeriesTool())
|
||
|
||
# 创建包含数值和时间列的数据画像
|
||
profile = DataProfile(
|
||
file_path='test.csv',
|
||
row_count=100,
|
||
column_count=2,
|
||
columns=[
|
||
ColumnInfo(name='value', dtype='numeric', missing_rate=0.0, unique_count=50),
|
||
ColumnInfo(name='date', dtype='datetime', missing_rate=0.0, unique_count=100)
|
||
],
|
||
inferred_type='unknown'
|
||
)
|
||
|
||
applicable = registry.get_applicable_tools(profile)
|
||
|
||
# 所有工具都应该适用（GetColumnDistributionTool 适用于所有数据）
|
||
assert len(applicable) > 0
|
||
|
||
|
||
|
||
class TestToolManager:
|
||
"""测试工具管理器。"""
|
||
|
||
def test_select_tools_for_datetime_data(self):
|
||
"""测试为包含时间字段的数据选择工具。"""
|
||
from src.tools.tool_manager import ToolManager
|
||
|
||
# 创建工具注册表并注册所有工具
|
||
registry = ToolRegistry()
|
||
registry.register(GetTimeSeriesTool())
|
||
registry.register(CalculateTrendTool())
|
||
registry.register(GetColumnDistributionTool())
|
||
|
||
manager = ToolManager(registry)
|
||
|
||
# 创建包含时间字段的数据画像
|
||
profile = DataProfile(
|
||
file_path='test.csv',
|
||
row_count=100,
|
||
column_count=1,
|
||
columns=[
|
||
ColumnInfo(name='date', dtype='datetime', missing_rate=0.0, unique_count=100)
|
||
],
|
||
inferred_type='unknown',
|
||
key_fields={},
|
||
quality_score=100.0,
|
||
summary='Test data'
|
||
)
|
||
|
||
tools = manager.select_tools(profile)
|
||
tool_names = [tool.name for tool in tools]
|
||
|
||
# 应该包含时间序列工具
|
||
assert 'get_time_series' in tool_names
|
||
assert 'calculate_trend' in tool_names
|
||
|
||
def test_select_tools_for_numeric_data(self):
|
||
"""测试为包含数值字段的数据选择工具。"""
|
||
from src.tools.tool_manager import ToolManager
|
||
|
||
registry = ToolRegistry()
|
||
registry.register(CalculateStatisticsTool())
|
||
registry.register(DetectOutliersTool())
|
||
registry.register(GetCorrelationTool())
|
||
|
||
manager = ToolManager(registry)
|
||
|
||
# 创建包含数值字段的数据画像
|
||
profile = DataProfile(
|
||
file_path='test.csv',
|
||
row_count=100,
|
||
column_count=2,
|
||
columns=[
|
||
ColumnInfo(name='value1', dtype='numeric', missing_rate=0.0, unique_count=50),
|
||
ColumnInfo(name='value2', dtype='numeric', missing_rate=0.0, unique_count=50)
|
||
],
|
||
inferred_type='unknown',
|
||
key_fields={},
|
||
quality_score=100.0,
|
||
summary='Test data'
|
||
)
|
||
|
||
tools = manager.select_tools(profile)
|
||
tool_names = [tool.name for tool in tools]
|
||
|
||
# 应该包含统计工具
|
||
assert 'calculate_statistics' in tool_names
|
||
assert 'detect_outliers' in tool_names
|
||
assert 'get_correlation' in tool_names
|
||
|
||
def test_select_tools_for_categorical_data(self):
|
||
"""测试为包含分类字段的数据选择工具。"""
|
||
from src.tools.tool_manager import ToolManager
|
||
|
||
registry = ToolRegistry()
|
||
registry.register(GetColumnDistributionTool())
|
||
registry.register(GetValueCountsTool())
|
||
registry.register(PerformGroupbyTool())
|
||
|
||
manager = ToolManager(registry)
|
||
|
||
# 创建包含分类字段的数据画像
|
||
profile = DataProfile(
|
||
file_path='test.csv',
|
||
row_count=100,
|
||
column_count=1,
|
||
columns=[
|
||
ColumnInfo(name='category', dtype='categorical', missing_rate=0.0, unique_count=5)
|
||
],
|
||
inferred_type='unknown',
|
||
key_fields={},
|
||
quality_score=100.0,
|
||
summary='Test data'
|
||
)
|
||
|
||
tools = manager.select_tools(profile)
|
||
tool_names = [tool.name for tool in tools]
|
||
|
||
# 应该包含分类工具
|
||
assert 'get_column_distribution' in tool_names
|
||
assert 'get_value_counts' in tool_names
|
||
assert 'perform_groupby' in tool_names
|
||
|
||
def test_no_geo_tools_for_non_geo_data(self):
|
||
"""测试不为非地理数据选择地理工具。"""
|
||
from src.tools.tool_manager import ToolManager
|
||
|
||
registry = ToolRegistry()
|
||
registry.register(GetColumnDistributionTool())
|
||
|
||
manager = ToolManager(registry)
|
||
|
||
# 创建不包含地理字段的数据画像
|
||
profile = DataProfile(
|
||
file_path='test.csv',
|
||
row_count=100,
|
||
column_count=1,
|
||
columns=[
|
||
ColumnInfo(name='value', dtype='numeric', missing_rate=0.0, unique_count=50)
|
||
],
|
||
inferred_type='unknown',
|
||
key_fields={},
|
||
quality_score=100.0,
|
||
summary='Test data'
|
||
)
|
||
|
||
tools = manager.select_tools(profile)
|
||
tool_names = [tool.name for tool in tools]
|
||
|
||
# 不应该包含地理工具
|
||
assert 'create_map_visualization' not in tool_names
|
||
|
||
def test_identify_missing_tools(self):
|
||
"""测试识别缺失的工具。"""
|
||
from src.tools.tool_manager import ToolManager
|
||
|
||
# 创建空的工具注册表
|
||
empty_registry = ToolRegistry()
|
||
manager = ToolManager(empty_registry)
|
||
|
||
# 创建包含时间字段的数据画像
|
||
profile = DataProfile(
|
||
file_path='test.csv',
|
||
row_count=100,
|
||
column_count=1,
|
||
columns=[
|
||
ColumnInfo(name='date', dtype='datetime', missing_rate=0.0, unique_count=100)
|
||
],
|
||
inferred_type='unknown',
|
||
key_fields={},
|
||
quality_score=100.0,
|
||
summary='Test data'
|
||
)
|
||
|
||
# 尝试选择工具
|
||
tools = manager.select_tools(profile)
|
||
|
||
# 获取缺失的工具
|
||
missing = manager.get_missing_tools()
|
||
|
||
# 应该识别出缺失的时间序列工具
|
||
assert len(missing) > 0
|
||
assert any(tool in missing for tool in ['get_time_series', 'calculate_trend'])
|
||
|
||
def test_clear_missing_tools(self):
|
||
"""测试清空缺失工具列表。"""
|
||
from src.tools.tool_manager import ToolManager
|
||
|
||
empty_registry = ToolRegistry()
|
||
manager = ToolManager(empty_registry)
|
||
|
||
# 创建数据画像并选择工具（会记录缺失工具）
|
||
profile = DataProfile(
|
||
file_path='test.csv',
|
||
row_count=100,
|
||
column_count=1,
|
||
columns=[
|
||
ColumnInfo(name='date', dtype='datetime', missing_rate=0.0, unique_count=100)
|
||
],
|
||
inferred_type='unknown',
|
||
key_fields={},
|
||
quality_score=100.0,
|
||
summary='Test data'
|
||
)
|
||
|
||
manager.select_tools(profile)
|
||
assert len(manager.get_missing_tools()) > 0
|
||
|
||
# 清空缺失工具列表
|
||
manager.clear_missing_tools()
|
||
assert len(manager.get_missing_tools()) == 0
|
||
|
||
def test_get_tool_descriptions(self):
|
||
"""测试获取工具描述。"""
|
||
from src.tools.tool_manager import ToolManager
|
||
|
||
registry = ToolRegistry()
|
||
tool1 = GetColumnDistributionTool()
|
||
tool2 = CalculateStatisticsTool()
|
||
registry.register(tool1)
|
||
registry.register(tool2)
|
||
|
||
manager = ToolManager(registry)
|
||
|
||
tools = [tool1, tool2]
|
||
descriptions = manager.get_tool_descriptions(tools)
|
||
|
||
assert len(descriptions) == 2
|
||
assert all('name' in desc for desc in descriptions)
|
||
assert all('description' in desc for desc in descriptions)
|
||
assert all('parameters' in desc for desc in descriptions)
|
||
|
||
def test_tool_deduplication(self):
|
||
"""测试工具去重。"""
|
||
from src.tools.tool_manager import ToolManager
|
||
|
||
registry = ToolRegistry()
|
||
# 注册一个工具，它可能被多个类别选中
|
||
tool = GetColumnDistributionTool()
|
||
registry.register(tool)
|
||
|
||
manager = ToolManager(registry)
|
||
|
||
# 创建包含多种类型字段的数据画像
|
||
profile = DataProfile(
|
||
file_path='test.csv',
|
||
row_count=100,
|
||
column_count=2,
|
||
columns=[
|
||
ColumnInfo(name='category', dtype='categorical', missing_rate=0.0, unique_count=5),
|
||
ColumnInfo(name='value', dtype='numeric', missing_rate=0.0, unique_count=50)
|
||
],
|
||
inferred_type='unknown',
|
||
key_fields={},
|
||
quality_score=100.0,
|
||
summary='Test data'
|
||
)
|
||
|
||
tools = manager.select_tools(profile)
|
||
tool_names = [tool.name for tool in tools]
|
||
|
||
# 工具名称应该是唯一的（没有重复）
|
||
assert len(tool_names) == len(set(tool_names))
|