wehub-resource-sync
101 lines
3.1 KiB
Python
101 lines
3.1 KiB
Python
"""Tests for ClientManager connection lifecycle (no MongoDB instance required).
|
|
|
|
These tests verify that ClientManager properly closes the underlying
|
|
AsyncMongoClient when all references are released, and keeps it alive
|
|
while references remain.
|
|
"""
|
|
|
|
import pytest
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
pytest.importorskip(
|
|
"pymongo",
|
|
reason="pymongo is required for Mongo ClientManager tests",
|
|
)
|
|
|
|
from lightrag.kg.mongo_impl import ClientManager
|
|
|
|
|
|
class TestClientManagerLifecycle:
|
|
"""Verify ClientManager connection open/close behavior."""
|
|
|
|
def _reset_manager(self):
|
|
"""Reset ClientManager class state between tests."""
|
|
ClientManager._instances = {"client": None, "db": None, "ref_count": 0}
|
|
|
|
def teardown_method(self):
|
|
self._reset_manager()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_release_client_closes_connection_when_ref_count_zero(self):
|
|
"""When ref_count drops to 0, the MongoClient should be closed and cleared."""
|
|
mock_client = AsyncMock()
|
|
mock_db = MagicMock()
|
|
|
|
ClientManager._instances = {
|
|
"client": mock_client,
|
|
"db": mock_db,
|
|
"ref_count": 1,
|
|
}
|
|
|
|
await ClientManager.release_client(mock_db)
|
|
|
|
mock_client.close.assert_awaited_once()
|
|
assert ClientManager._instances["client"] is None
|
|
assert ClientManager._instances["db"] is None
|
|
assert ClientManager._instances["ref_count"] == 0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_release_client_keeps_connection_with_multiple_refs(self):
|
|
"""When other references exist, the MongoClient must NOT be closed."""
|
|
mock_client = AsyncMock()
|
|
mock_db = MagicMock()
|
|
|
|
ClientManager._instances = {
|
|
"client": mock_client,
|
|
"db": mock_db,
|
|
"ref_count": 3,
|
|
}
|
|
|
|
await ClientManager.release_client(mock_db)
|
|
|
|
mock_client.close.assert_not_awaited()
|
|
assert ClientManager._instances["ref_count"] == 2
|
|
assert ClientManager._instances["client"] is mock_client
|
|
assert ClientManager._instances["db"] is mock_db
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_release_client_noop_for_wrong_db(self):
|
|
"""Releasing a db that is not the tracked instance should do nothing."""
|
|
mock_client = AsyncMock()
|
|
mock_db = MagicMock()
|
|
other_db = MagicMock()
|
|
|
|
ClientManager._instances = {
|
|
"client": mock_client,
|
|
"db": mock_db,
|
|
"ref_count": 1,
|
|
}
|
|
|
|
await ClientManager.release_client(other_db)
|
|
|
|
mock_client.close.assert_not_awaited()
|
|
assert ClientManager._instances["ref_count"] == 1
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_release_client_noop_for_none(self):
|
|
"""Releasing None should be a safe no-op."""
|
|
mock_client = AsyncMock()
|
|
mock_db = MagicMock()
|
|
|
|
ClientManager._instances = {
|
|
"client": mock_client,
|
|
"db": mock_db,
|
|
"ref_count": 1,
|
|
}
|
|
|
|
await ClientManager.release_client(None)
|
|
|
|
mock_client.close.assert_not_awaited()
|
|
assert ClientManager._instances["ref_count"] == 1
|