265 lines
7.5 KiB
Python
265 lines
7.5 KiB
Python
"""
|
||
TraceStudio 执行引擎和缓存系统集成测试
|
||
"""
|
||
import sys
|
||
import asyncio
|
||
from pathlib import Path
|
||
import time
|
||
|
||
# 添加项目根目录到路径
|
||
project_root = Path(__file__).parent.parent.parent
|
||
sys.path.insert(0, str(project_root))
|
||
|
||
from server.app.core.cache_manager import CacheManager, MemoryCache, DiskCache, CacheEvictionPolicy
|
||
from server.app.core.workflow_executor import WorkflowExecutor, WorkflowGraph
|
||
from server.app.nodes.example_nodes import AddNode, MultiplyNode
|
||
|
||
|
||
# ============= 测试 1: 内存缓存 =============
|
||
|
||
def test_memory_cache():
|
||
"""测试内存缓存"""
|
||
print("\n" + "="*60)
|
||
print("测试 1: 内存缓存")
|
||
print("="*60)
|
||
|
||
cache = MemoryCache(max_size=3, ttl=None, policy=CacheEvictionPolicy.LRU)
|
||
|
||
# 添加数据
|
||
cache.set("key1", {"value": 1})
|
||
cache.set("key2", {"value": 2})
|
||
cache.set("key3", {"value": 3})
|
||
print("✅ 添加 3 个缓存项")
|
||
|
||
# 获取数据
|
||
assert cache.get("key1") == {"value": 1}
|
||
print("✅ 成功获取 key1")
|
||
|
||
# 触发淘汰（LRU）
|
||
cache.set("key4", {"value": 4})
|
||
print("✅ 添加第 4 个项，触发 LRU 淘汰")
|
||
|
||
# key2 应该被淘汰（最久未使用）
|
||
assert cache.get("key2") is None
|
||
print("✅ key2 被淘汰（LRU）")
|
||
|
||
# 统计
|
||
stats = cache.get_stats()
|
||
print(f"\n缓存统计: {stats}")
|
||
|
||
|
||
def test_disk_cache():
|
||
"""测试磁盘缓存"""
|
||
print("\n" + "="*60)
|
||
print("测试 2: 磁盘缓存")
|
||
print("="*60)
|
||
|
||
cache_dir = Path("/tmp/tracestudio_cache_test")
|
||
cache = DiskCache(cache_dir, ttl=None)
|
||
|
||
# 设置数据
|
||
cache.set("test_key", {"data": "test_value", "count": 42})
|
||
print("✅ 设置磁盘缓存")
|
||
|
||
# 获取数据
|
||
result = cache.get("test_key")
|
||
assert result == {"data": "test_value", "count": 42}
|
||
print("✅ 成功获取磁盘缓存数据")
|
||
|
||
# 统计
|
||
stats = cache.get_stats()
|
||
print(f"\n缓存统计: {stats}")
|
||
|
||
# 清理
|
||
cache.clear()
|
||
print("✅ 清理磁盘缓存")
|
||
|
||
|
||
def test_cache_ttl():
|
||
"""测试缓存 TTL"""
|
||
print("\n" + "="*60)
|
||
print("测试 3: 缓存 TTL")
|
||
print("="*60)
|
||
|
||
cache = MemoryCache(max_size=100, ttl=1) # 1 秒过期
|
||
|
||
# 设置数据
|
||
cache.set("expire_key", {"value": "will_expire"})
|
||
print("✅ 设置缓存（1秒过期）")
|
||
|
||
# 立即获取
|
||
assert cache.get("expire_key") == {"value": "will_expire"}
|
||
print("✅ 立即获取成功")
|
||
|
||
# 等待过期
|
||
print("⏳ 等待 1.1 秒...")
|
||
time.sleep(1.1)
|
||
|
||
# 再次获取应该返回 None
|
||
assert cache.get("expire_key") is None
|
||
print("✅ 过期后返回 None")
|
||
|
||
|
||
# ============= 测试 4: 工作流图 =============
|
||
|
||
def test_workflow_graph():
|
||
"""测试工作流图"""
|
||
print("\n" + "="*60)
|
||
print("测试 4: 工作流图")
|
||
print("="*60)
|
||
|
||
graph = WorkflowGraph()
|
||
|
||
# 添加节点
|
||
# 按照执行器约定：type 使用 NodeType 值（例如 'normal'），实现类名放到 class_name
|
||
graph.add_node("n1", "normal", {"offset": 0})
|
||
graph.nodes["n1"]["class_name"] = "AddNode"
|
||
graph.add_node("n2", "normal", {"scale": 2.0})
|
||
graph.nodes["n2"]["class_name"] = "MultiplyNode"
|
||
graph.add_node("n3", "normal", {"offset": 10})
|
||
graph.nodes["n3"]["class_name"] = "AddNode"
|
||
print("✅ 添加 3 个节点")
|
||
|
||
# 添加连接
|
||
graph.add_edge("n1", "result", "n2", "a")
|
||
graph.add_edge("n2", "result", "n3", "a")
|
||
print("✅ 添加连接: n1 -> n2 -> n3")
|
||
|
||
# 获取依赖
|
||
deps = graph.get_dependencies("n3")
|
||
assert "n1" in deps and "n2" in deps
|
||
print(f"✅ n3 的依赖: {deps}")
|
||
|
||
# 拓扑排序
|
||
order = graph.topological_sort()
|
||
assert order == ["n1", "n2", "n3"]
|
||
print(f"✅ 拓扑排序: {order}")
|
||
|
||
# 检查循环
|
||
assert not graph.has_cycle()
|
||
print("✅ 无循环依赖")
|
||
|
||
|
||
# ============= 测试 5: 工作流执行 =============
|
||
|
||
async def test_workflow_execution():
|
||
"""测试工作流执行"""
|
||
print("\n" + "="*60)
|
||
print("测试 5: 工作流执行")
|
||
print("="*60)
|
||
|
||
executor = WorkflowExecutor(user_id="test_user")
|
||
|
||
# 定义工作流
|
||
nodes = [
|
||
{
|
||
"id": "add_node",
|
||
"type": "AddNode",
|
||
"params": {"offset": 5}
|
||
},
|
||
{
|
||
"id": "multiply_node",
|
||
"type": "MultiplyNode",
|
||
"params": {"scale": 2.0}
|
||
}
|
||
]
|
||
|
||
|
||
# 边定义：add_node.result -> multiply_node.a
|
||
edges = [
|
||
{
|
||
"source": "add_node",
|
||
"source_port": "result",
|
||
"target": "multiply_node",
|
||
"target_port": "a"
|
||
}
|
||
]
|
||
|
||
# add_node 的 a、b，multiply_node 的 b 都由全局 context 提供
|
||
global_context = {"a": 3, "b": 4}
|
||
|
||
try:
|
||
success, report = await executor.execute(nodes, edges, global_context=global_context)
|
||
|
||
if success:
|
||
print(f"✅ 工作流执行成功")
|
||
print(f" 执行 ID: {report.get('execution_id')}")
|
||
print(f" 总耗时: {report.get('total_duration', 0):.3f}s")
|
||
print(f" 节点数: {len(report.get('nodes', {}))}")
|
||
for node_id, info in report.get('nodes', {}).items():
|
||
print(f" - {node_id}: {info.get('status')}")
|
||
else:
|
||
print(f"❌ 工作流执行失败: {report.get('error') if report else 'Unknown error'}")
|
||
if report:
|
||
print(f" 详情: {report}")
|
||
|
||
except Exception as e:
|
||
print(f"❌ 执行异常: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
|
||
|
||
# ============= 测试 6: 缓存集成 =============
|
||
|
||
def test_cache_manager():
|
||
"""测试缓存管理器"""
|
||
print("\n" + "="*60)
|
||
print("测试 6: 缓存管理器")
|
||
print("="*60)
|
||
|
||
# 初始化缓存
|
||
CacheManager.init_memory_cache(max_size=100, ttl=None)
|
||
CacheManager.init_disk_cache(Path("/tmp/tracestudio_cache"), ttl=None)
|
||
print("✅ 初始化内存和磁盘缓存")
|
||
|
||
# 设置数据
|
||
CacheManager.set("test_key", {"data": "value"}, storage="memory")
|
||
CacheManager.set("test_key2", {"data": "value2"}, storage="disk")
|
||
print("✅ 分别设置内存和磁盘缓存")
|
||
|
||
# 获取数据
|
||
assert CacheManager.get("test_key", prefer="memory") == {"data": "value"}
|
||
assert CacheManager.get("test_key2", prefer="disk") == {"data": "value2"}
|
||
print("✅ 成功获取数据")
|
||
|
||
# 同时设置
|
||
CacheManager.set_both("both_key", {"data": "both"})
|
||
assert CacheManager.get("both_key") == {"data": "both"}
|
||
print("✅ 同时设置内存和磁盘缓存成功")
|
||
|
||
# 统计
|
||
stats = CacheManager.get_stats()
|
||
print(f"\n缓存统计:")
|
||
if stats.get("memory"):
|
||
print(f" 内存: {stats['memory']}")
|
||
if stats.get("disk"):
|
||
print(f" 磁盘: {stats['disk']}")
|
||
|
||
# 清理
|
||
CacheManager.clear()
|
||
print("✅ 清理所有缓存")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
print("\n" + "🧪 TraceStudio 执行引擎和缓存系统测试".center(60, "="))
|
||
|
||
try:
|
||
# 运行测试
|
||
test_memory_cache()
|
||
test_disk_cache()
|
||
test_cache_ttl()
|
||
test_workflow_graph()
|
||
asyncio.run(test_workflow_execution())
|
||
test_cache_manager()
|
||
|
||
print("\n" + "="*60)
|
||
print("✅ 所有测试通过！")
|
||
print("="*60 + "\n")
|
||
|
||
except Exception as e:
|
||
print("\n" + "="*60)
|
||
print(f"❌ 测试失败: {e}")
|
||
print("="*60 + "\n")
|
||
import traceback
|
||
traceback.print_exc()
|