90 lines
3.0 KiB
Python
90 lines
3.0 KiB
Python
"""性能优化相关测试"""
|
|
import pytest
|
|
import sys
|
|
from pathlib import Path
|
|
sys.path.insert(0, str(Path(__file__).parent.parent / 'src'))
|
|
|
|
from graph.sqlite_storage import GraphSQLiteStorage, GraphNodeRecord, GraphEdgeRecord
|
|
from storage.embedding_cache import EmbeddingCache, text_hash
|
|
from datetime import datetime
|
|
|
|
|
|
class TestGraphSQLiteStorage:
|
|
"""图存储测试"""
|
|
|
|
def test_upsert_and_get_node(self, tmp_path):
|
|
storage = GraphSQLiteStorage(db_path=str(tmp_path / "test_graph.db"))
|
|
|
|
node = GraphNodeRecord(
|
|
id="node1",
|
|
entity="牧尘",
|
|
node_type="person",
|
|
properties={"role": "developer"},
|
|
importance=0.9,
|
|
created_at=datetime.now()
|
|
)
|
|
storage.upsert_node(node)
|
|
|
|
result = storage.get_node("node1")
|
|
assert result is not None
|
|
assert result.entity == "牧尘"
|
|
assert result.properties["role"] == "developer"
|
|
|
|
def test_upsert_and_get_edge(self, tmp_path):
|
|
storage = GraphSQLiteStorage(db_path=str(tmp_path / "test_graph2.db"))
|
|
|
|
node1 = GraphNodeRecord("n1", "A", "entity", {}, 1.0, datetime.now())
|
|
node2 = GraphNodeRecord("n2", "B", "entity", {}, 1.0, datetime.now())
|
|
storage.upsert_node(node1)
|
|
storage.upsert_node(node2)
|
|
|
|
edge = GraphEdgeRecord(
|
|
id="e1", source_id="n1", target_id="n2",
|
|
relation="depends_on", properties={},
|
|
created_at=datetime.now()
|
|
)
|
|
storage.upsert_edge(edge)
|
|
|
|
edges = storage.get_edges("n1", direction="out")
|
|
assert len(edges) == 1
|
|
assert edges[0].relation == "depends_on"
|
|
|
|
def test_count(self, tmp_path):
|
|
storage = GraphSQLiteStorage(db_path=str(tmp_path / "test_graph3.db"))
|
|
node = GraphNodeRecord("n1", "X", "entity", {}, 1.0, datetime.now())
|
|
storage.upsert_node(node)
|
|
assert storage.count_nodes() == 1
|
|
|
|
|
|
class TestEmbeddingCache:
|
|
"""Embedding 缓存测试"""
|
|
|
|
def test_cache_hit(self, tmp_path):
|
|
cache = EmbeddingCache(db_path=str(tmp_path / "test_emb.db"))
|
|
|
|
vector = [0.1, 0.2, 0.3]
|
|
cache.set("测试文本", vector)
|
|
|
|
result = cache.get("测试文本")
|
|
assert result == vector
|
|
|
|
def test_cache_miss(self, tmp_path):
|
|
cache = EmbeddingCache(db_path=str(tmp_path / "test_emb2.db"))
|
|
|
|
result = cache.get("不存在的文本")
|
|
assert result is None
|
|
|
|
def test_text_hash(self):
|
|
h = text_hash("hello")
|
|
assert len(h) == 16
|
|
assert text_hash("hello") == text_hash("hello") # 确定性
|
|
assert text_hash("hello") != text_hash("world")
|
|
|
|
def test_stats(self, tmp_path):
|
|
cache = EmbeddingCache(db_path=str(tmp_path / "test_emb3.db"))
|
|
cache.set("文本1", [0.1, 0.2])
|
|
cache.set("文本2", [0.3, 0.4])
|
|
|
|
stats = cache.stats()
|
|
assert stats["count"] == 2
|
|
assert stats["size_bytes"] > 0 |