File size: 3,963 Bytes
979853c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 | import importlib
from unittest.mock import AsyncMock, patch
import pytest
import lightrag.utils as utils_module
from lightrag.kg.postgres_impl import PGGraphStorage, PostgreSQLDB
from lightrag.namespace import NameSpace
def make_db() -> PostgreSQLDB:
return PostgreSQLDB(
{
"host": "localhost",
"port": 5432,
"user": "postgres",
"password": "postgres",
"database": "postgres",
"workspace": "test_ws",
"max_connections": 10,
"connection_retry_attempts": 3,
"connection_retry_backoff": 0,
"connection_retry_backoff_max": 0,
"pool_close_timeout": 5.0,
}
)
@pytest.mark.asyncio
async def test_execute_timing_logs_success():
db = make_db()
async def fake_run_with_retry(operation, **kwargs):
conn = AsyncMock()
conn.execute = AsyncMock(return_value="INSERT 0 1")
await operation(conn)
db._run_with_retry = AsyncMock(side_effect=fake_run_with_retry)
with patch("lightrag.kg.postgres_impl.performance_timing_log") as timing_log:
await db.execute("SELECT 1", timing_label="test label")
assert any(
"connection.execute completed" in call.args[0]
for call in timing_log.call_args_list
)
@pytest.mark.asyncio
async def test_execute_timing_logs_failure():
db = make_db()
async def fake_run_with_retry(operation, **kwargs):
conn = AsyncMock()
conn.execute = AsyncMock(side_effect=RuntimeError("boom"))
await operation(conn)
db._run_with_retry = AsyncMock(side_effect=fake_run_with_retry)
with patch("lightrag.kg.postgres_impl.performance_timing_log") as timing_log:
with pytest.raises(RuntimeError, match="boom"):
await db.execute("SELECT 1", timing_label="test label")
assert any(
"connection.execute failed" in call.args[0]
for call in timing_log.call_args_list
)
@pytest.mark.asyncio
async def test_graph_upsert_node_passes_timing_label():
storage = PGGraphStorage(
namespace=NameSpace.GRAPH_STORE_CHUNK_ENTITY_RELATION,
workspace="test_ws",
global_config={},
embedding_func=AsyncMock(),
)
storage.graph_name = "test_graph"
storage._query = AsyncMock(return_value=[])
await storage.upsert_node(
"node-1",
{
"entity_id": "node-1",
"description": "desc",
},
)
assert storage._query.await_args.kwargs["timing_label"] == (
"test_ws PGGraphStorage.upsert_node"
)
@pytest.mark.asyncio
async def test_graph_upsert_edge_passes_timing_label():
storage = PGGraphStorage(
namespace=NameSpace.GRAPH_STORE_CHUNK_ENTITY_RELATION,
workspace="test_ws",
global_config={},
embedding_func=AsyncMock(),
)
storage.graph_name = "test_graph"
storage._query = AsyncMock(return_value=[])
await storage.upsert_edge(
"node-1",
"node-2",
{
"weight": 1.0,
"description": "desc",
},
)
assert storage._query.await_args.kwargs["timing_label"] == (
"test_ws PGGraphStorage.upsert_edge"
)
def test_performance_timing_logs_reads_new_env_only(monkeypatch):
with monkeypatch.context() as m:
m.setenv("LIGHTRAG_DOC_QUERY_TIMING_LOGS", "false")
m.setenv("LIGHTRAG_PERFORMANCE_TIMING_LOGS", "true")
reloaded = importlib.reload(utils_module)
assert reloaded.PERFORMANCE_TIMING_LOGS is True
importlib.reload(utils_module)
def test_performance_timing_logs_ignores_old_env(monkeypatch):
with monkeypatch.context() as m:
m.setenv("LIGHTRAG_DOC_QUERY_TIMING_LOGS", "true")
m.setenv("LIGHTRAG_PERFORMANCE_TIMING_LOGS", "false")
reloaded = importlib.reload(utils_module)
assert reloaded.PERFORMANCE_TIMING_LOGS is False
importlib.reload(utils_module)
|