248 lines
8.9 KiB
Python
248 lines
8.9 KiB
Python
"""
|
||
Tests for graph module.
|
||
"""
|
||
import pytest
|
||
import tempfile
|
||
from pathlib import Path
|
||
from datetime import datetime
|
||
import networkx as nx
|
||
|
||
from src.graph.models import GraphNode, GraphEdge, NodeType, EdgeType
|
||
from src.graph.builder import GraphBuilder
|
||
from src.graph.query import GraphQuery
|
||
from src.graph.indexer import LazyIndexer
|
||
from src.models.distilled import Distilled
|
||
|
||
|
||
class TestGraphNode:
|
||
def test_create_entity_node(self):
|
||
node = GraphNode.from_entity(
|
||
entity="牧尘",
|
||
distilled_id="test-123",
|
||
episode_id="ep-456",
|
||
confidence=0.8,
|
||
importance=2
|
||
)
|
||
assert node.type == NodeType.ENTITY
|
||
assert node.label == "牧尘"
|
||
assert node.distilled_id == "test-123"
|
||
|
||
def test_create_fact_node(self):
|
||
node = GraphNode.from_fact(
|
||
fact="牧尘喜欢简洁的回答",
|
||
distilled_id="test-123",
|
||
episode_id="ep-456",
|
||
confidence=0.9
|
||
)
|
||
assert node.type == NodeType.FACT
|
||
assert "简洁" in node.label
|
||
|
||
def test_to_dict_roundtrip(self):
|
||
node = GraphNode.from_entity("测试", "d1", "e1")
|
||
d = node.to_dict()
|
||
node2 = GraphNode.from_dict(d)
|
||
assert node.id == node2.id
|
||
assert node.label == node2.label
|
||
assert node.type == node2.type
|
||
|
||
def test_node_importance_assignment(self):
|
||
node = GraphNode.from_entity("重要人物", "d1", "e1", importance=3)
|
||
assert node.importance == 3
|
||
|
||
|
||
class TestGraphEdge:
|
||
def test_create_edge(self):
|
||
node_a = GraphNode.from_entity("A", "d1", "e1")
|
||
node_b = GraphNode.from_entity("B", "d2", "e2")
|
||
edge = GraphEdge.create_relation(
|
||
source_id=node_a.id,
|
||
target_id=node_b.id,
|
||
relation=EdgeType.RELATES_TO,
|
||
weight=0.9
|
||
)
|
||
assert edge.source == node_a.id
|
||
assert edge.target == node_b.id
|
||
assert edge.weight == 0.9
|
||
|
||
def test_to_dict_roundtrip(self):
|
||
node_a = GraphNode.from_entity("A", "d1", "e1")
|
||
node_b = GraphNode.from_entity("B", "d2", "e2")
|
||
edge = GraphEdge.create_relation(node_a.id, node_b.id, EdgeType.SIMILAR_TO, weight=0.8)
|
||
d = edge.to_dict()
|
||
edge2 = GraphEdge.from_dict(d)
|
||
assert edge.id == edge2.id
|
||
assert edge.relation == edge2.relation
|
||
|
||
|
||
class TestGraphBuilder:
|
||
def test_build_graph_from_distilled(self):
|
||
"""Test building a graph from a Distilled record."""
|
||
distilled = Distilled(
|
||
id="distilled-1",
|
||
episode_id="ep-1",
|
||
type="fact",
|
||
summary="测试总结",
|
||
facts=[
|
||
"牧尘喜欢简洁",
|
||
"牧尘偏好Python",
|
||
"牧尘用Arch Linux"
|
||
],
|
||
entities=["牧尘"],
|
||
importance=2
|
||
)
|
||
|
||
builder = GraphBuilder()
|
||
builder.add_distilled(distilled)
|
||
graph = builder.graph
|
||
|
||
# Should create entity nodes + fact nodes
|
||
assert len(graph.nodes) >= 4 # 至少：entity "牧尘" + 3 facts
|
||
|
||
def test_entity_extraction_from_facts(self):
|
||
"""Test that entities are properly extracted from facts."""
|
||
distilled = Distilled(
|
||
id="distilled-2",
|
||
episode_id="ep-2",
|
||
type="fact",
|
||
summary="",
|
||
facts=["牧尘使用ComfyUI"],
|
||
entities=["牧尘", "ComfyUI"],
|
||
importance=0.7
|
||
)
|
||
|
||
builder = GraphBuilder()
|
||
builder.add_distilled(distilled)
|
||
graph = builder.graph
|
||
|
||
# Both entity and fact should create nodes
|
||
assert len(graph.nodes) >= 2
|
||
|
||
|
||
class TestGraphQuery:
|
||
def setup_method(self):
|
||
"""Set up a test graph using NetworkX directly."""
|
||
# Create a simple graph: A -> B -> C, A -> D
|
||
self.graph = nx.DiGraph()
|
||
|
||
self.node_a = GraphNode.from_entity("A", "d1", "e1")
|
||
self.node_b = GraphNode.from_entity("B", "d2", "e2")
|
||
self.node_c = GraphNode.from_entity("C", "d3", "e3")
|
||
self.node_d = GraphNode.from_entity("D", "d4", "e4")
|
||
|
||
# Add nodes to graph with data
|
||
self.graph.add_node(self.node_a.id, **self.node_a.to_dict())
|
||
self.graph.add_node(self.node_b.id, **self.node_b.to_dict())
|
||
self.graph.add_node(self.node_c.id, **self.node_c.to_dict())
|
||
self.graph.add_node(self.node_d.id, **self.node_d.to_dict())
|
||
|
||
# A -> B -> C (2 hops), A -> D (1 hop)
|
||
self.graph.add_edge(self.node_a.id, self.node_b.id, **GraphEdge.create_relation(
|
||
self.node_a.id, self.node_b.id, EdgeType.RELATES_TO, weight=0.9
|
||
).to_dict())
|
||
self.graph.add_edge(self.node_b.id, self.node_c.id, **GraphEdge.create_relation(
|
||
self.node_b.id, self.node_c.id, EdgeType.RELATES_TO, weight=0.9
|
||
).to_dict())
|
||
self.graph.add_edge(self.node_a.id, self.node_d.id, **GraphEdge.create_relation(
|
||
self.node_a.id, self.node_d.id, EdgeType.RELATES_TO, weight=0.8
|
||
).to_dict())
|
||
|
||
def test_bfs_1_hop(self):
|
||
"""Test BFS with 1 hop."""
|
||
query = GraphQuery(self.graph)
|
||
result = query.bfs_n_hops(self.node_a.id, max_hops=1)
|
||
|
||
# 1 hop should reach B and D (not C which is 2 hops away)
|
||
assert 1 in result
|
||
found_1hop = result.get(1, set())
|
||
assert self.node_b.id in found_1hop
|
||
assert self.node_d.id in found_1hop
|
||
assert self.node_c.id not in found_1hop
|
||
|
||
def test_bfs_2_hops(self):
|
||
"""Test BFS with 2 hops."""
|
||
query = GraphQuery(self.graph)
|
||
result = query.bfs_n_hops(self.node_a.id, max_hops=2)
|
||
|
||
# 2 hops should reach C (through B)
|
||
assert 2 in result
|
||
found_2hops = result.get(2, set())
|
||
assert self.node_c.id in found_2hops
|
||
|
||
def test_degree_centrality(self):
|
||
"""Test degree centrality calculation."""
|
||
query = GraphQuery(self.graph)
|
||
centrality = query.degree_centrality()
|
||
|
||
# Graph has 4 nodes, n-1 = 3
|
||
# A: out-degree=2, degree=2 → 2/3
|
||
# B: out-degree=1, degree=2 → 2/3 (in+out)
|
||
# C: out-degree=0, degree=1 → 1/3
|
||
# D: out-degree=0, degree=1 → 1/3
|
||
assert centrality[self.node_a.id] == 2.0 / 3.0
|
||
assert centrality[self.node_b.id] == 2.0 / 3.0
|
||
assert centrality[self.node_c.id] == 1.0 / 3.0
|
||
assert centrality[self.node_d.id] == 1.0 / 3.0
|
||
|
||
def test_jaccard_similarity(self):
|
||
"""Test Jaccard similarity between nodes."""
|
||
# Add shared neighbor for A and B
|
||
node_e = GraphNode.from_entity("E", "d5", "e5")
|
||
self.graph.add_node(node_e.id, **node_e.to_dict())
|
||
|
||
# A -> E, B -> E (both connect to E)
|
||
self.graph.add_edge(self.node_a.id, node_e.id, **GraphEdge.create_relation(
|
||
self.node_a.id, node_e.id, EdgeType.RELATES_TO, weight=0.9
|
||
).to_dict())
|
||
self.graph.add_edge(self.node_b.id, node_e.id, **GraphEdge.create_relation(
|
||
self.node_b.id, node_e.id, EdgeType.RELATES_TO, weight=0.9
|
||
).to_dict())
|
||
|
||
query = GraphQuery(self.graph)
|
||
sim = query.jaccard_similarity(self.node_a.id, self.node_b.id)
|
||
|
||
# A has neighbors {B, D, E}, B has neighbors {C, E}
|
||
# Jaccard = |{E}| / |{B, D, E, C}| = 1/4
|
||
assert 0 < sim <= 1
|
||
|
||
|
||
class TestLazyIndexer:
|
||
def test_indexer_creation(self):
|
||
"""Test lazy indexer initialization."""
|
||
with tempfile.TemporaryDirectory() as tmpdir:
|
||
indexer = LazyIndexer(storage_dir=Path(tmpdir))
|
||
assert indexer.storage_dir.exists()
|
||
|
||
def test_build_and_lookup(self):
|
||
"""Test building indexes and looking up."""
|
||
with tempfile.TemporaryDirectory() as tmpdir:
|
||
# Create a simple graph
|
||
graph = nx.DiGraph()
|
||
node = GraphNode.from_entity("测试节点", "d1", "e1")
|
||
graph.add_node(node.id, **node.to_dict())
|
||
|
||
indexer = LazyIndexer(storage_dir=Path(tmpdir))
|
||
indexer.build_indexes(graph)
|
||
|
||
# Should find the node by label
|
||
results = indexer.lookup_by_label("测试节点")
|
||
assert len(results) == 1
|
||
|
||
def test_save_and_load_index(self):
|
||
"""Test saving and loading indexes."""
|
||
with tempfile.TemporaryDirectory() as tmpdir:
|
||
# Create a graph
|
||
graph = nx.DiGraph()
|
||
node = GraphNode.from_entity("测试节点", "d1", "e1")
|
||
graph.add_node(node.id, **node.to_dict())
|
||
|
||
# Build and save
|
||
indexer1 = LazyIndexer(storage_dir=Path(tmpdir))
|
||
indexer1.build_indexes(graph)
|
||
indexer1.save_index()
|
||
|
||
# Load in new indexer
|
||
indexer2 = LazyIndexer(storage_dir=Path(tmpdir))
|
||
indexer2.load_index()
|
||
|
||
results = indexer2.lookup_by_label("测试节点")
|
||
assert len(results) == 1 |