chore: import upstream snapshot with attribution

2026-07-13 12:08:54 +08:00
commit 4a4a1fed67
721 changed files with 262090 additions and 0 deletions
@@ -0,0 +1,138 @@
+import importlib
+from unittest.mock import AsyncMock, MagicMock, 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"
+    # upsert_edge drives the lock + cypher via db._run_with_retry, not _query.
+    storage.db = MagicMock()
+    storage.db._run_with_retry = AsyncMock(return_value=None)
+
+    await storage.upsert_edge(
+        "node-1",
+        "node-2",
+        {
+            "weight": 1.0,
+            "description": "desc",
+        },
+    )
+
+    assert storage.db._run_with_retry.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)