LightRAG/tests/test_unified_lock_safety.py

"""
Tests for UnifiedLock safety when lock is None.

This test module verifies that UnifiedLock raises RuntimeError instead of
allowing unprotected execution when the underlying lock is None, preventing
false security and potential race conditions.

Critical Bug 1: When self._lock is None, __aenter__ used to log WARNING but
still return successfully, allowing critical sections to run without lock
protection, causing race conditions and data corruption.

Critical Bug 2: In __aexit__, when async_lock.release() fails, the error
recovery logic would attempt to release it again, causing double-release issues.
"""

import pytest
from unittest.mock import MagicMock, AsyncMock
from lightrag.kg.shared_storage import UnifiedLock


class TestUnifiedLockSafety:
    """Test suite for UnifiedLock None safety checks."""

    @pytest.mark.asyncio
    async def test_unified_lock_raises_on_none_async(self):
        """
        Test that UnifiedLock raises RuntimeError when lock is None (async mode).

        Scenario: Attempt to use UnifiedLock before initialize_share_data() is called.
        Expected: RuntimeError raised, preventing unprotected critical section execution.
        """
        lock = UnifiedLock(
            lock=None, is_async=True, name="test_async_lock", enable_logging=False
        )

        with pytest.raises(
            RuntimeError, match="shared data not initialized|Lock.*is None"
        ):
            async with lock:
                # This code should NEVER execute
                pytest.fail(
                    "Code inside lock context should not execute when lock is None"
                )

    @pytest.mark.asyncio
    async def test_unified_lock_raises_on_none_sync(self):
        """
        Test that UnifiedLock raises RuntimeError when lock is None (sync mode).

        Scenario: Attempt to use UnifiedLock with None lock in sync mode.
        Expected: RuntimeError raised with clear error message.
        """
        lock = UnifiedLock(
            lock=None, is_async=False, name="test_sync_lock", enable_logging=False
        )

        with pytest.raises(
            RuntimeError, match="shared data not initialized|Lock.*is None"
        ):
            async with lock:
                # This code should NEVER execute
                pytest.fail(
                    "Code inside lock context should not execute when lock is None"
                )

    @pytest.mark.asyncio
    async def test_error_message_clarity(self):
        """
        Test that the error message clearly indicates the problem and solution.

        Scenario: Lock is None and user tries to acquire it.
        Expected: Error message mentions 'shared data not initialized' and
                  'initialize_share_data()'.
        """
        lock = UnifiedLock(
            lock=None,
            is_async=True,
            name="test_error_message",
            enable_logging=False,
        )

        with pytest.raises(RuntimeError) as exc_info:
            async with lock:
                pass

        error_message = str(exc_info.value)
        # Verify error message contains helpful information
        assert (
            "shared data not initialized" in error_message.lower()
            or "lock" in error_message.lower()
        )
        assert "initialize_share_data" in error_message or "None" in error_message

    @pytest.mark.asyncio
    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()

    @pytest.mark.asyncio
    async def test_aexit_recovery_on_main_lock_failure(self):
        """
        Test that __aexit__ recovery logic works when main lock release fails.

        Scenario: main_lock.release() fails before async_lock is attempted.
        Expected: Recovery logic should attempt to release async_lock to prevent
                  resource leaks.

        This verifies the recovery logic still works correctly with async_lock_released tracking.
        """
        # Create mock locks
        main_lock = MagicMock()
        main_lock.acquire = MagicMock()

        # Make main_lock.release() fail
        def mock_main_release_fail():
            raise RuntimeError("Main lock release failed")

        main_lock.release = MagicMock(side_effect=mock_main_release_fail)

        async_lock = AsyncMock()
        async_lock.acquire = AsyncMock()
        async_lock.release = MagicMock()

        # Create UnifiedLock with both locks (sync mode with async_lock)
        lock = UnifiedLock(
            lock=main_lock, is_async=False, name="test_recovery", 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 main lock release error
            assert "Main lock release failed" in str(e)

        # Main lock release should have been attempted
        main_lock.release.assert_called_once()

        # Recovery logic should have attempted to release async_lock
        async_lock.release.assert_called_once()