123 lines
4.2 KiB
Python
123 lines
4.2 KiB
Python
"""
|
|
Tests for UnifiedLock safety when lock is None.
|
|
|
|
This test module verifies that get_internal_lock() and get_data_init_lock()
|
|
raise RuntimeError when shared data is not initialized, preventing false
|
|
security and potential race conditions.
|
|
|
|
Design: The None check has been moved from UnifiedLock.__aenter__/__enter__
|
|
to the lock factory functions (get_internal_lock, get_data_init_lock) for
|
|
early failure detection.
|
|
|
|
Critical Bug 1 (Fixed): When self._lock is None, the code would fail with
|
|
AttributeError. Now the check is in factory functions for clearer errors.
|
|
|
|
Critical Bug 2: In __aexit__, when async_lock.release() fails, the error
|
|
recovery logic would attempt to release it again, causing double-release issues.
|
|
"""
|
|
|
|
from unittest.mock import MagicMock, AsyncMock
|
|
|
|
import pytest
|
|
|
|
from lightrag.kg.shared_storage import (
|
|
UnifiedLock,
|
|
get_internal_lock,
|
|
get_data_init_lock,
|
|
finalize_share_data,
|
|
)
|
|
|
|
|
|
class TestUnifiedLockSafety:
|
|
"""Test suite for UnifiedLock None safety checks."""
|
|
|
|
def setup_method(self):
|
|
"""Ensure shared data is finalized before each test."""
|
|
finalize_share_data()
|
|
|
|
def teardown_method(self):
|
|
"""Clean up after each test."""
|
|
finalize_share_data()
|
|
|
|
def test_get_internal_lock_raises_when_not_initialized(self):
|
|
"""
|
|
Test that get_internal_lock() raises RuntimeError when shared data is not initialized.
|
|
|
|
Scenario: Call get_internal_lock() before initialize_share_data() is called.
|
|
Expected: RuntimeError raised with clear error message.
|
|
|
|
This test verifies the None check has been moved to the factory function.
|
|
"""
|
|
with pytest.raises(
|
|
RuntimeError, match="Shared data not initialized.*initialize_share_data"
|
|
):
|
|
get_internal_lock()
|
|
|
|
def test_get_data_init_lock_raises_when_not_initialized(self):
|
|
"""
|
|
Test that get_data_init_lock() raises RuntimeError when shared data is not initialized.
|
|
|
|
Scenario: Call get_data_init_lock() before initialize_share_data() is called.
|
|
Expected: RuntimeError raised with clear error message.
|
|
|
|
This test verifies the None check has been moved to the factory function.
|
|
"""
|
|
with pytest.raises(
|
|
RuntimeError, match="Shared data not initialized.*initialize_share_data"
|
|
):
|
|
get_data_init_lock()
|
|
|
|
@pytest.mark.offline
|
|
async def test_aexit_no_double_release_on_async_lock_failure(self):
|
|
"""
|
|
Test that __aexit__ doesn't attempt to release async_lock twice when it fails.
|
|
|
|
Scenario: async_lock.release() fails during normal release.
|
|
Expected: Recovery logic should NOT attempt to release async_lock again,
|
|
preventing double-release issues.
|
|
|
|
This tests Bug 2 fix: async_lock_released tracking prevents double release.
|
|
"""
|
|
# Create mock locks
|
|
main_lock = MagicMock()
|
|
main_lock.acquire = MagicMock()
|
|
main_lock.release = MagicMock()
|
|
|
|
async_lock = AsyncMock()
|
|
async_lock.acquire = AsyncMock()
|
|
|
|
# Make async_lock.release() fail
|
|
release_call_count = 0
|
|
|
|
def mock_release_fail():
|
|
nonlocal release_call_count
|
|
release_call_count += 1
|
|
raise RuntimeError("Async lock release failed")
|
|
|
|
async_lock.release = MagicMock(side_effect=mock_release_fail)
|
|
|
|
# Create UnifiedLock with both locks (sync mode with async_lock)
|
|
lock = UnifiedLock(
|
|
lock=main_lock,
|
|
is_async=False,
|
|
name="test_double_release",
|
|
enable_logging=False,
|
|
)
|
|
lock._async_lock = async_lock
|
|
|
|
# Try to use the lock - should fail during __aexit__
|
|
try:
|
|
async with lock:
|
|
pass
|
|
except RuntimeError as e:
|
|
# Should get the async lock release error
|
|
assert "Async lock release failed" in str(e)
|
|
|
|
# Verify async_lock.release() was called only ONCE, not twice
|
|
assert (
|
|
release_call_count == 1
|
|
), f"async_lock.release() should be called only once, but was called {release_call_count} times"
|
|
|
|
# Main lock should have been released successfully
|
|
main_lock.release.assert_called_once()
|